Merge pull request #22964 from perfinion:bazelrc
PiperOrigin-RevId: 225413451
diff --git a/.bazelrc b/.bazelrc
index 8d9834f..cd7e13d 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -76,7 +76,6 @@
build --define=use_fast_cpp_protos=true
build --define=allow_oversize_protos=true
-build --define=grpc_no_ares=true
build --spawn_strategy=standalone
build --genrule_strategy=standalone
diff --git a/CODEOWNERS b/CODEOWNERS
index 54a61a4..cb3fa23 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -1,7 +1,7 @@
# Where component owners are known, add them here.
/tenosrflow/core/debug @caisq
-/tensorflow/core/nccl/ @azaks @csigg
+/tensorflow/core/nccl/ @azaks2 @chsigg
/tensorflow/core/platform/windows/ @mrry
/tensorflow/core/platform/s3 @yongtang
/tensorflow/go @asimshankar
@@ -51,13 +51,13 @@
/tensorflow/contrib/pi_examples/ @maciekcc
/tensorflow/contrib/quantization/ @petewarden
/tensorflow/contrib/rnn/ @ebrevdo @scottzhu
-/tensorflow/contrib/saved_model/ @nfiedel @sukritiramesh @allenl
+/tensorflow/contrib/saved_model/ @nfiedel @sukritiramesh @allenlavoie
/tensorflow/contrib/seq2seq/ @ebrevdo @lmthang
/tensorflow/contrib/session_bundle/ @nfiedel @sukritiramesh
/tensorflow/contrib/slim/ @sguada @thenbasilmanran
/tensorflow/contrib/stateless/ @girving @alextp
/tensorflow/contrib/tensor_forest/ @gilberthendry @thomascolthurst @yupbank
-/tensorflow/contrib/tensorrt/ @aaroey
+/tensorflow/contrib/tensorrt/ @aaroey @smit-hinsu @azaks2
# NEED OWNER: /tensorflow/contrib/testing/
/tensorflow/contrib/timeseries/ @allenlavoie
/tensorflow/contrib/tpu/ @frankchn @saeta @jhseu @sourabhbajaj
diff --git a/README.md b/README.md
index 0441749..68d7e18 100644
--- a/README.md
+++ b/README.md
@@ -113,9 +113,10 @@
Build Type | Status | Artifacts
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------
**IBM s390x** | [](http://ibmz-ci.osuosl.org/job/TensorFlow_IBMZ_CI/) | TBA
-**IBM ppc64le CPU** | [](http://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Build/) | TBA
-**IBM ppc64le GPU** Nightly | [](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/) | [Nightly](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/)
-**IBM ppc64le GPU** Stable Release | [](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/) | [Release](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/)
+**Linux ppc64le CPU** Nightly | [](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Build/) | [Nightly](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Nightly_Artifact/)
+**Linux ppc64le CPU** Stable Release | [](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Release_Build/) | [Release](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Release_Build/)
+**Linux ppc64le GPU** Nightly | [](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Build/) | [Nightly](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/)
+**Linux ppc64le GPU** Stable Release | [](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/) | [Release](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/)
**Linux CPU with Intel® MKL-DNN** Nightly | [](https://tensorflow-ci.intel.com/job/tensorflow-mkl-linux-cpu/) | [Nightly](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-whl-nightly/)
**Linux CPU with Intel® MKL-DNN** Python 2.7<br> **Linux CPU with Intel® MKL-DNN** Python 3.4<br> **Linux CPU with Intel® MKL-DNN** Python 3.5<br> **Linux CPU with Intel® MKL-DNN** Python 3.6 | [](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-release-whl/lastStableBuild) | [1.11.0 py2.7](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp27-cp27mu-linux_x86_64.whl)<br>[1.11.0 py3.4](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp34-cp34m-linux_x86_64.whl)<br>[1.11.0 py3.5](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp35-cp35m-linux_x86_64.whl)<br>[1.11.0 py3.6](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp36-cp36m-linux_x86_64.whl)
diff --git a/RELEASE.md b/RELEASE.md
index b13b071..32abdce 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -7,6 +7,8 @@
Serving.
* Keras models now support evaluating with a `tf.data.Dataset`.
* TensorFlow binaries are built with XLA support linked in by default.
+* Ignite Dataset added to contrib/ignite that allows to work with Apache
+ Ignite.
## Bug Fixes and Other Changes
diff --git a/WORKSPACE b/WORKSPACE
index 0f59c44..7057d3f 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -16,30 +16,27 @@
closure_repositories()
-http_archive(
- name = "base_images_docker",
- sha256 = "e2b1b7254270bb7605e814a9dbf6d1e4ae04a11136ff1714fbfdabe3f87f7cf9",
- strip_prefix = "base-images-docker-12801524f867e657fbb5d1a74f31618aff181ac6",
- urls = ["https://github.com/GoogleCloudPlatform/base-images-docker/archive/12801524f867e657fbb5d1a74f31618aff181ac6.tar.gz"],
+load("//third_party/toolchains/preconfig/generate:archives.bzl",
+ "bazel_toolchains_archive")
+
+bazel_toolchains_archive()
+
+load(
+ "@bazel_toolchains//repositories:repositories.bzl",
+ bazel_toolchains_repositories = "repositories",
)
-http_archive(
- name = "bazel_toolchains",
- sha256 = "15b5858b1b5541ec44df31b94c3b8672815b31d71215a98398761ea9f4c4eedb",
- strip_prefix = "bazel-toolchains-6200b238c9c2d137c0d9a7262c80cc71d98e692b",
- urls = [
- "https://github.com/bazelbuild/bazel-toolchains/archive/6200b238c9c2d137c0d9a7262c80cc71d98e692b.tar.gz",
- ],
+bazel_toolchains_repositories()
+
+load(
+ "@io_bazel_rules_docker//container:container.bzl",
+ container_repositories = "repositories",
)
-http_archive(
- name = "io_bazel_rules_docker",
- sha256 = "29d109605e0d6f9c892584f07275b8c9260803bf0c6fcb7de2623b2bedc910bd",
- strip_prefix = "rules_docker-0.5.1",
- urls = ["https://github.com/bazelbuild/rules_docker/archive/v0.5.1.tar.gz"],
-)
+container_repositories()
-load("//third_party/toolchains/preconfig/generate:workspace.bzl", "remote_config_workspace")
+load("//third_party/toolchains/preconfig/generate:workspace.bzl",
+ "remote_config_workspace")
remote_config_workspace()
diff --git a/configure.py b/configure.py
index d19607a..7f05bc6 100644
--- a/configure.py
+++ b/configure.py
@@ -477,7 +477,8 @@
print('Please upgrade your bazel installation to version %s or higher to '
'build TensorFlow!' % min_version)
sys.exit(1)
- if curr_version_int > max_version_int and not 'TF_IGNORE_MAX_BAZEL_VERSION' in os.environ:
+ if (curr_version_int > max_version_int and
+ 'TF_IGNORE_MAX_BAZEL_VERSION' not in os.environ):
print('Please downgrade your bazel installation to version %s or lower to '
'build TensorFlow!' % max_version)
sys.exit(1)
diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index fd4b942..823ad8f 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -606,9 +606,11 @@
name = "tensorflow_py",
srcs_version = "PY2AND3",
visibility = ["//visibility:public"],
- deps = [
+ deps = select({
+ "api_version_2": [],
+ "//conditions:default": ["//tensorflow/contrib:contrib_py"],
+ }) + [
":tensorflow_py_no_contrib",
- "//tensorflow/contrib:contrib_py",
"//tensorflow/python/estimator:estimator_py",
],
)
diff --git a/tensorflow/api_template.__init__.py b/tensorflow/api_template.__init__.py
index d81cf06..4eba763 100644
--- a/tensorflow/api_template.__init__.py
+++ b/tensorflow/api_template.__init__.py
@@ -20,14 +20,14 @@
import os as _os
+# API IMPORTS PLACEHOLDER
+
# pylint: disable=g-bad-import-order
from tensorflow.python.tools import component_api_helper as _component_api_helper
_component_api_helper.package_hook(
parent_package_str=__name__,
child_package_str=('tensorflow_estimator.python.estimator.api.estimator'))
-# API IMPORTS PLACEHOLDER
-
# Make sure directory containing top level submodules is in
# the __path__ so that "from tensorflow.foo import bar" works.
# We're using bitwise, but there's nothing special about that.
diff --git a/tensorflow/api_template_v1.__init__.py b/tensorflow/api_template_v1.__init__.py
index 65bdb6c..21b5277 100644
--- a/tensorflow/api_template_v1.__init__.py
+++ b/tensorflow/api_template_v1.__init__.py
@@ -23,13 +23,13 @@
# pylint: disable=g-bad-import-order
from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import
+# API IMPORTS PLACEHOLDER
+
from tensorflow.python.tools import component_api_helper as _component_api_helper
_component_api_helper.package_hook(
parent_package_str=__name__,
child_package_str=('tensorflow_estimator.python.estimator.api.estimator'))
-# API IMPORTS PLACEHOLDER
-
from tensorflow.python.util.lazy_loader import LazyLoader # pylint: disable=g-import-not-at-top
contrib = LazyLoader('contrib', globals(), 'tensorflow.contrib')
del LazyLoader
diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD
index 59c23e7..25df970 100644
--- a/tensorflow/c/BUILD
+++ b/tensorflow/c/BUILD
@@ -176,6 +176,34 @@
)
tf_cuda_library(
+ name = "env",
+ srcs = [
+ "env.cc",
+ ],
+ hdrs = [
+ "env.h",
+ ],
+ copts = tf_copts(),
+ visibility = ["//visibility:public"],
+ deps = select({
+ "//tensorflow:android": [
+ ":c_api",
+ ":tf_status_helper",
+ "//tensorflow/core:android_tensorflow_lib_lite",
+ "//tensorflow/core:platform_env",
+ "//tensorflow/core:lib",
+ ],
+ "//conditions:default": [
+ ":c_api",
+ ":tf_status_helper",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:platform_env",
+ "//tensorflow/core:lib",
+ ],
+ }) + [":c_api_internal"],
+)
+
+tf_cuda_library(
name = "kernels",
srcs = [
"kernels.cc",
@@ -189,11 +217,13 @@
"//tensorflow:android": [
":c_api",
":c_api_internal",
+ ":tf_status_helper",
"//tensorflow/core:android_tensorflow_lib_lite",
],
"//conditions:default": [
":c_api",
":c_api_internal",
+ ":tf_status_helper",
"//tensorflow/core:framework",
],
}),
@@ -333,6 +363,27 @@
)
tf_cuda_cc_test(
+ name = "env_test",
+ size = "small",
+ srcs = ["env_test.cc"],
+ linkopts = select({
+ "//tensorflow:darwin": ["-headerpad_max_install_names"],
+ "//conditions:default": [],
+ }),
+ tags = ["noasan"],
+ # We must ensure that the dependencies can be dynamically linked since
+ # the shared library must be able to use core:framework.
+ # linkstatic = tf_kernel_tests_linkstatic(),
+ deps = [
+ ":c_api",
+ ":env",
+ "//tensorflow/core:lib",
+ "//tensorflow/core:test",
+ "//tensorflow/core:test_main",
+ ],
+)
+
+tf_cuda_cc_test(
name = "kernels_test",
size = "small",
srcs = ["kernels_test.cc"],
diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc
index 94d18eb..9580215 100644
--- a/tensorflow/c/c_api.cc
+++ b/tensorflow/c/c_api.cc
@@ -488,6 +488,7 @@
// Non-static for testing.
TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src,
TF_Status* status) {
+ TF_SetStatus(status, TF_OK, "");
if (!src.IsInitialized()) {
status->status = FailedPrecondition(
"attempt to use a tensor with an uninitialized value");
diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h
index 3d56268..c7abba8 100644
--- a/tensorflow/c/c_api.h
+++ b/tensorflow/c/c_api.h
@@ -91,7 +91,7 @@
// --------------------------------------------------------------------------
// TF_Version returns a string describing version information of the
// TensorFlow library. TensorFlow using semantic versioning.
-TF_CAPI_EXPORT extern const char* TF_Version();
+TF_CAPI_EXPORT extern const char* TF_Version(void);
// --------------------------------------------------------------------------
// TF_DataType holds the type for a scalar value. E.g., one slot in a tensor.
@@ -157,7 +157,7 @@
typedef struct TF_Status TF_Status;
// Return a new status object.
-TF_CAPI_EXPORT extern TF_Status* TF_NewStatus();
+TF_CAPI_EXPORT extern TF_Status* TF_NewStatus(void);
// Delete a previously created status object.
TF_CAPI_EXPORT extern void TF_DeleteStatus(TF_Status*);
@@ -196,7 +196,7 @@
size_t proto_len);
// Useful for passing *out* a protobuf.
-TF_CAPI_EXPORT extern TF_Buffer* TF_NewBuffer();
+TF_CAPI_EXPORT extern TF_Buffer* TF_NewBuffer(void);
TF_CAPI_EXPORT extern void TF_DeleteBuffer(TF_Buffer*);
@@ -305,7 +305,7 @@
typedef struct TF_SessionOptions TF_SessionOptions;
// Return a new options object.
-TF_CAPI_EXPORT extern TF_SessionOptions* TF_NewSessionOptions();
+TF_CAPI_EXPORT extern TF_SessionOptions* TF_NewSessionOptions(void);
// Set the target in TF_SessionOptions.options.
// target can be empty, a single entry, or a comma separated list of entries.
@@ -338,7 +338,7 @@
typedef struct TF_Graph TF_Graph;
// Return a new graph object.
-TF_CAPI_EXPORT extern TF_Graph* TF_NewGraph();
+TF_CAPI_EXPORT extern TF_Graph* TF_NewGraph(void);
// Destroy an options object. Graph will be deleted once no more
// TFSession's are referencing it.
@@ -890,7 +890,8 @@
// TF_GraphImportGraphDef.
typedef struct TF_ImportGraphDefOptions TF_ImportGraphDefOptions;
-TF_CAPI_EXPORT extern TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions();
+TF_CAPI_EXPORT extern TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions(
+ void);
TF_CAPI_EXPORT extern void TF_DeleteImportGraphDefOptions(
TF_ImportGraphDefOptions* opts);
@@ -1611,7 +1612,7 @@
//
// The data in the buffer will be the serialized OpList proto for ops registered
// in this address space.
-TF_CAPI_EXPORT extern TF_Buffer* TF_GetAllOpList();
+TF_CAPI_EXPORT extern TF_Buffer* TF_GetAllOpList(void);
// TF_ApiDefMap encapsulates a collection of API definitions for an operation.
//
diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h
index 80c8bfe..3e3a485 100644
--- a/tensorflow/c/c_api_experimental.h
+++ b/tensorflow/c/c_api_experimental.h
@@ -239,7 +239,7 @@
// Platform-specific implementation to return an unused port. (This should used
// in tests only.)
-TF_CAPI_EXPORT int TF_PickUnusedPortOrDie();
+TF_CAPI_EXPORT int TF_PickUnusedPortOrDie(void);
// Fast path method that makes constructing a single scalar tensor require less
// overhead and copies.
diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h
index 8d6c8d9..120748a 100755
--- a/tensorflow/c/eager/c_api.h
+++ b/tensorflow/c/eager/c_api.h
@@ -48,7 +48,7 @@
typedef struct TFE_ContextOptions TFE_ContextOptions;
// Return a new options object.
-TF_CAPI_EXPORT extern TFE_ContextOptions* TFE_NewContextOptions();
+TF_CAPI_EXPORT extern TFE_ContextOptions* TFE_NewContextOptions(void);
// Set the config in TF_ContextOptions.options.
// config should be a serialized tensorflow.ConfigProto proto.
@@ -170,23 +170,11 @@
int dim_index,
TF_Status* status);
-// Returns the device of the operation that produced `h`.
-// If `h` was produced by a copy, returns the destination device of
-// the copy. Note that returned device name is not always the device
-// holding the tensor handle's memory. If you want the latter, use
-// TFE_TensorHandleBackingDeviceName.
-// This function will block till the operation that produces `h` has completed.
-//
-// Device on which the kernel of the operation that produced `h` ran.
-//
-// If `h` was produced by a copy, returns the destination device of
-// the copy.
-//
-// Note that returned device name is not always the device that owns the memory
-// that backs the tensor handle. For the latter see
-// TFE_TensorHandleBackingDeviceName.
-//
-// This function will block till the operation that produces `h` has completed.
+// Returns the device of the operation that produced `h`. If `h` was produced by
+// a copy, returns the destination device of the copy. Note that the returned
+// device name is not always the device holding the tensor handle's memory. If
+// you want the latter, use TFE_TensorHandleBackingDeviceName. This function
+// will block till the operation that produces `h` has completed.
TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName(
TFE_TensorHandle* h, TF_Status* status);
diff --git a/tensorflow/c/env.cc b/tensorflow/c/env.cc
new file mode 100644
index 0000000..1c35ff9
--- /dev/null
+++ b/tensorflow/c/env.cc
@@ -0,0 +1,183 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/c/env.h"
+
+#include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/types.h"
+
+struct TF_StringStream {
+ std::vector<::tensorflow::string>* list;
+ size_t position;
+};
+
+void TF_CreateDir(const char* dirname, TF_Status* status) {
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(
+ status, ::tensorflow::Env::Default()->CreateDir(dirname));
+}
+
+void TF_DeleteDir(const char* dirname, TF_Status* status) {
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(
+ status, ::tensorflow::Env::Default()->DeleteDir(dirname));
+}
+
+void TF_DeleteRecursively(const char* dirname, int64_t* undeleted_file_count,
+ int64_t* undeleted_dir_count, TF_Status* status) {
+ ::tensorflow::int64 f, d;
+
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(
+ status, ::tensorflow::Env::Default()->DeleteRecursively(dirname, &f, &d));
+ *undeleted_file_count = f;
+ *undeleted_dir_count = d;
+}
+
+void TF_FileStat(const char* filename, TF_FileStatistics* stats,
+ TF_Status* status) {
+ ::tensorflow::FileStatistics cc_stats;
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Status s =
+ ::tensorflow::Env::Default()->Stat(filename, &cc_stats);
+ ::tensorflow::Set_TF_Status_from_Status(status, s);
+ if (s.ok()) {
+ stats->length = cc_stats.length;
+ stats->mtime_nsec = cc_stats.mtime_nsec;
+ stats->is_directory = cc_stats.is_directory;
+ }
+}
+
+void TF_NewWritableFile(const char* filename, TF_WritableFileHandle** handle,
+ TF_Status* status) {
+ std::unique_ptr<::tensorflow::WritableFile> f;
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Status s =
+ ::tensorflow::Env::Default()->NewWritableFile(filename, &f);
+ ::tensorflow::Set_TF_Status_from_Status(status, s);
+
+ if (s.ok()) {
+ *handle = reinterpret_cast<TF_WritableFileHandle*>(f.release());
+ }
+}
+
+void TF_CloseWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
+ auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Close());
+ delete cc_file;
+}
+
+void TF_SyncWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
+ auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Sync());
+}
+
+void TF_FlushWritableFile(TF_WritableFileHandle* handle, TF_Status* status) {
+ auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Flush());
+}
+
+void TF_AppendWritableFile(TF_WritableFileHandle* handle, const char* data,
+ size_t length, TF_Status* status) {
+ auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle);
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(
+ status, cc_file->Append(::tensorflow::StringPiece{data, length}));
+}
+
+void TF_DeleteFile(const char* filename, TF_Status* status) {
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(
+ status, ::tensorflow::Env::Default()->DeleteFile(filename));
+}
+
+bool TF_StringStreamNext(TF_StringStream* list, const char** result) {
+ if (list->position >= list->list->size()) {
+ *result = nullptr;
+ return false;
+ }
+
+ *result = list->list->at(list->position++).c_str();
+ return true;
+}
+
+void TF_StringStreamDone(TF_StringStream* list) {
+ delete list->list;
+ delete list;
+}
+TF_StringStream* TF_GetChildren(const char* dirname, TF_Status* status) {
+ auto* children = new std::vector<::tensorflow::string>;
+
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(
+ status, ::tensorflow::Env::Default()->GetChildren(dirname, children));
+
+ auto* list = new TF_StringStream;
+ list->list = children;
+ list->position = 0;
+ return list;
+}
+
+TF_StringStream* TF_GetLocalTempDirectories() {
+ auto* tmpdirs = new std::vector<::tensorflow::string>;
+
+ ::tensorflow::Env::Default()->GetLocalTempDirectories(tmpdirs);
+
+ auto* list = new TF_StringStream;
+ list->list = tmpdirs;
+ list->position = 0;
+ return list;
+}
+
+TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void) {
+ return ::tensorflow::Env::Default()->NowNanos();
+}
+
+// Returns the number of microseconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowMicros(void) {
+ return ::tensorflow::Env::Default()->NowMicros();
+}
+
+// Returns the number of seconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowSeconds(void) {
+ return ::tensorflow::Env::Default()->NowSeconds();
+}
+
+void TF_DefaultThreadOptions(TF_ThreadOptions* options) {
+ options->stack_size = 0;
+ options->guard_size = 0;
+ options->numa_node = -1;
+}
+
+TF_Thread* TF_StartThread(const TF_ThreadOptions* options,
+ const char* thread_name, void (*work_func)(void*),
+ void* param) {
+ ::tensorflow::ThreadOptions cc_options;
+ cc_options.stack_size = options->stack_size;
+ cc_options.guard_size = options->guard_size;
+ cc_options.numa_node = options->numa_node;
+ return reinterpret_cast<TF_Thread*>(::tensorflow::Env::Default()->StartThread(
+ cc_options, thread_name, [=]() { (*work_func)(param); }));
+}
+
+void TF_JoinThread(TF_Thread* thread) {
+ // ::tensorflow::Thread joins on destruction
+ delete reinterpret_cast<::tensorflow::Thread*>(thread);
+}
diff --git a/tensorflow/c/env.h b/tensorflow/c/env.h
new file mode 100644
index 0000000..1565235
--- /dev/null
+++ b/tensorflow/c/env.h
@@ -0,0 +1,194 @@
+/* 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.
+==============================================================================*/
+
+#include <stddef.h>
+#include <stdint.h>
+
+#ifndef TENSORFLOW_C_ENV_H_
+#define TENSORFLOW_C_ENV_H_
+
+#include "tensorflow/c/c_api.h"
+
+// --------------------------------------------------------------------------
+// C API for tensorflow::Env.
+
+struct TF_WritableFileHandle;
+struct TF_StringStream;
+struct TF_Thread;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct TF_FileStatistics {
+ // The length of the file in bytes.
+ int64_t length;
+ // The last modified time in nanoseconds.
+ int64_t mtime_nsec;
+ // Whether the name refers to a directory.
+ bool is_directory;
+} TF_FileStatistics;
+
+typedef struct TF_ThreadOptions {
+ // Thread stack size to use (in bytes), zero implies that the system default
+ // will be used.
+ size_t stack_size;
+
+ // Guard area size to use near thread stacks to use (in bytes), zero implies
+ // that the system default will be used.
+ size_t guard_size;
+
+ // The NUMA node to use, -1 implies that there should be no NUMA affinity for
+ // this thread.
+ int numa_node;
+} TF_ThreadOptions;
+
+// Creates the specified directory. Typical status code are:
+// * TF_OK - successfully created the directory
+// * TF_ALREADY_EXISTS - directory already exists
+// * TF_PERMISSION_DENIED - dirname is not writable
+TF_CAPI_EXPORT extern void TF_CreateDir(const char* dirname, TF_Status* status);
+
+// Deletes the specified directory. Typical status codes are:
+// * TF_OK - successfully deleted the directory
+// * TF_FAILED_PRECONDITION - the directory is not empty
+TF_CAPI_EXPORT extern void TF_DeleteDir(const char* dirname, TF_Status* status);
+
+// Deletes the specified directory and all subdirectories and files underneath
+// it. This is accomplished by traversing the directory tree rooted at dirname
+// and deleting entries as they are encountered.
+//
+// If dirname itself is not readable or does not exist, *undeleted_dir_count is
+// set to 1, *undeleted_file_count is set to 0 and an appropriate status (e.g.
+// TF_NOT_FOUND) is returned.
+//
+// If dirname and all its descendants were successfully deleted, TF_OK is
+// returned and both error counters are set to zero.
+//
+// Otherwise, while traversing the tree, undeleted_file_count and
+// undeleted_dir_count are updated if an entry of the corresponding type could
+// not be deleted. The returned error status represents the reason that any one
+// of these entries could not be deleted.
+//
+// Typical status codes:
+// * TF_OK - dirname exists and we were able to delete everything underneath
+// * TF_NOT_FOUND - dirname doesn't exist
+// * TF_PERMISSION_DENIED - dirname or some descendant is not writable
+// * TF_UNIMPLEMENTED - some underlying functions (like Delete) are not
+// implemented
+TF_CAPI_EXPORT extern void TF_DeleteRecursively(const char* dirname,
+ int64_t* undeleted_file_count,
+ int64_t* undeleted_dir_count,
+ TF_Status* status);
+
+// Obtains statistics for the given path. If status is TF_OK, *stats is
+// updated, otherwise it is not touched.
+TF_CAPI_EXPORT extern void TF_FileStat(const char* filename,
+ TF_FileStatistics* stats,
+ TF_Status* status);
+
+// Creates or truncates the given filename and returns a handle to be used for
+// appending data to the file. If status is TF_OK, *handle is updated and the
+// caller is responsible for freeing it (see TF_CloseWritableFile).
+TF_CAPI_EXPORT extern void TF_NewWritableFile(const char* filename,
+ TF_WritableFileHandle** handle,
+ TF_Status* status);
+
+// Closes the given handle and frees its memory. If there was a problem closing
+// the file, it is indicated by status. Memory is freed in any case.
+TF_CAPI_EXPORT extern void TF_CloseWritableFile(TF_WritableFileHandle* handle,
+ TF_Status* status);
+
+// Syncs content of the handle to the filesystem. Blocks waiting for the
+// filesystem to indicate that the content has been persisted.
+TF_CAPI_EXPORT extern void TF_SyncWritableFile(TF_WritableFileHandle* handle,
+ TF_Status* status);
+
+// Flush local buffers to the filesystem. If the process terminates after a
+// successful flush, the contents may still be persisted, since the underlying
+// filesystem may eventually flush the contents. If the OS or machine crashes
+// after a successful flush, the contents may or may not be persisted, depending
+// on the implementation.
+TF_CAPI_EXPORT extern void TF_FlushWritableFile(TF_WritableFileHandle* handle,
+ TF_Status* status);
+
+// Appends the given bytes to the file. Any failure to do so is indicated in
+// status.
+TF_CAPI_EXPORT extern void TF_AppendWritableFile(TF_WritableFileHandle* handle,
+ const char* data,
+ size_t length,
+ TF_Status* status);
+
+// Deletes the named file and indicates whether successful in *status.
+TF_CAPI_EXPORT extern void TF_DeleteFile(const char* filename,
+ TF_Status* status);
+
+// Retrieves the next item from the given TF_StringStream and places a pointer
+// to it in *result. If no more items are in the list, *result is set to NULL
+// and false is returned.
+//
+// Ownership of the items retrieved with this function remains with the library.
+// Item points are invalidated after a call to TF_StringStreamDone.
+TF_CAPI_EXPORT extern bool TF_StringStreamNext(TF_StringStream* list,
+ const char** result);
+
+// Frees the resources associated with given string list. All pointers returned
+// by TF_StringStreamNext are invalid after this call.
+TF_CAPI_EXPORT extern void TF_StringStreamDone(TF_StringStream* list);
+
+// Retrieves the list of children of the given directory. You can iterate
+// through the list with TF_StringStreamNext. The caller is responsible for
+// freeing the list (see TF_StringStreamDone).
+TF_CAPI_EXPORT extern TF_StringStream* TF_GetChildren(const char* filename,
+ TF_Status* status);
+
+// Retrieves a list of directory names on the local machine that may be used for
+// temporary storage. You can iterate through the list with TF_StringStreamNext.
+// The caller is responsible for freeing the list (see TF_StringStreamDone).
+TF_CAPI_EXPORT extern TF_StringStream* TF_GetLocalTempDirectories(void);
+
+// Returns the number of nanoseconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void);
+
+// Returns the number of microseconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowMicros(void);
+
+// Returns the number of seconds since the Unix epoch.
+TF_CAPI_EXPORT extern uint64_t TF_NowSeconds(void);
+
+// Populates a TF_ThreadOptions struct with system-default values.
+TF_CAPI_EXPORT extern void TF_DefaultThreadOptions(TF_ThreadOptions* options);
+
+// Returns a new thread that is running work_func and is identified
+// (for debugging/performance-analysis) by thread_name.
+//
+// The given param (which may be null) is passed to work_func when the thread
+// starts. In this way, data may be passed from the thread back to the caller.
+//
+// Caller takes ownership of the result and must call TF_JoinThread on it
+// eventually.
+TF_CAPI_EXPORT extern TF_Thread* TF_StartThread(const TF_ThreadOptions* options,
+ const char* thread_name,
+ void (*work_func)(void*),
+ void* param);
+
+// Waits for the given thread to finish execution, then deletes it.
+TF_CAPI_EXPORT extern void TF_JoinThread(TF_Thread* thread);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // TENSORFLOW_C_ENV_H_
diff --git a/tensorflow/c/env_test.cc b/tensorflow/c/env_test.cc
new file mode 100644
index 0000000..687ad02
--- /dev/null
+++ b/tensorflow/c/env_test.cc
@@ -0,0 +1,127 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/c/env.h"
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+#define ASSERT_TF_OK(x) ASSERT_EQ(TF_OK, TF_GetCode(x))
+
+TEST(TestEnv, TestDirHandling) {
+ TF_StringStream* tempdirs = TF_GetLocalTempDirectories();
+ const char* tempdir;
+ bool found = false;
+ while (TF_StringStreamNext(tempdirs, &tempdir)) {
+ found = true;
+
+ TF_Status* s = TF_NewStatus();
+
+ ::tensorflow::string dirpath =
+ ::tensorflow::io::JoinPath(tempdir, "somedir");
+ TF_CreateDir(dirpath.c_str(), s);
+ ASSERT_TF_OK(s) << "TF_CreateDir failed for " << dirpath << ": "
+ << TF_Message(s);
+
+ ::tensorflow::string filepath =
+ ::tensorflow::io::JoinPath(dirpath, "somefile.txt");
+ TF_WritableFileHandle* handle;
+ TF_NewWritableFile(filepath.c_str(), &handle, s);
+ ASSERT_TF_OK(s) << "NewWritableFile failed for " << filepath << ": "
+ << TF_Message(s);
+
+ const char* data = "Hello, world!\n";
+ TF_AppendWritableFile(handle, data, strlen(data), s);
+ ASSERT_TF_OK(s) << "TF_AppendWritableFile failed to append data to file at "
+ << filepath << ": " << TF_Message(s);
+
+ TF_CloseWritableFile(handle, s);
+ ASSERT_TF_OK(s) << "TF_CloseWritableFile failed to close handle to "
+ << filepath << ": " << TF_Message(s);
+
+ TF_StringStream* children = TF_GetChildren(dirpath.c_str(), s);
+ ASSERT_TF_OK(s) << "TF_GetChildren failed for " << dirpath;
+ const char* childpath;
+ ASSERT_TRUE(TF_StringStreamNext(children, &childpath));
+ ASSERT_EQ(::tensorflow::string(childpath), "somefile.txt");
+ // There should only be one file in this directory.
+ ASSERT_FALSE(TF_StringStreamNext(children, &childpath));
+ ASSERT_EQ(childpath, nullptr);
+ TF_StringStreamDone(children);
+
+ TF_FileStatistics stats;
+ TF_FileStat(filepath.c_str(), &stats, s);
+ ASSERT_EQ(stats.length, strlen(data));
+ ASSERT_FALSE(stats.is_directory);
+ ASSERT_GT(stats.mtime_nsec, 0);
+
+ // Trying to delete a non-empty directory should fail.
+ TF_DeleteDir(dirpath.c_str(), s);
+ ASSERT_NE(TF_OK, TF_GetCode(s))
+ << "TF_DeleteDir unexpectedly succeeded with a non-empty directory "
+ << dirpath;
+
+ TF_DeleteFile(filepath.c_str(), s);
+ ASSERT_TF_OK(s) << "TF_DeleteFile failed for " << filepath << ": "
+ << TF_Message(s);
+
+ // Now deleting the directory should work.
+ TF_DeleteDir(dirpath.c_str(), s);
+ ASSERT_TF_OK(s) << "TF_DeleteDir failed for " << dirpath << ": "
+ << TF_Message(s);
+
+ TF_DeleteStatus(s);
+ break;
+ }
+
+ ASSERT_TRUE(found) << "expected at least one temp dir";
+
+ TF_StringStreamDone(tempdirs);
+}
+
+TEST(TestEnv, TestTimeFunctions) {
+ ASSERT_GE(TF_NowSeconds(), 946684800); // Midnight Jan 1, 2000
+ ASSERT_GE(TF_NowMicros(), 946684800 * 1e6);
+ ASSERT_GE(TF_NowNanos(), 946684800 * 1e9);
+}
+
+namespace {
+
+struct SomeThreadData {
+ ::tensorflow::mutex mu;
+ bool did_work = false;
+};
+
+void SomeThreadFunc(void* data) {
+ auto* real_data = static_cast<SomeThreadData*>(data);
+ ::tensorflow::mutex_lock l(real_data->mu);
+ real_data->did_work = true;
+}
+
+} // namespace
+
+TEST(TestEnv, TestThreads) {
+ TF_ThreadOptions options;
+ TF_DefaultThreadOptions(&options);
+ SomeThreadData data;
+ TF_Thread* thread =
+ TF_StartThread(&options, "SomeThreadName", &SomeThreadFunc, &data);
+ TF_JoinThread(thread);
+ ::tensorflow::mutex_lock l(data.mu);
+ ASSERT_TRUE(data.did_work);
+}
diff --git a/tensorflow/c/kernels.cc b/tensorflow/c/kernels.cc
index 3caa5bc..2a4eaec 100644
--- a/tensorflow/c/kernels.cc
+++ b/tensorflow/c/kernels.cc
@@ -17,6 +17,7 @@
#include "tensorflow/c/c_api_internal.h"
#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/framework/kernel_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
@@ -141,3 +142,19 @@
*tensor = result;
}
}
+
+void TF_SetOutput(TF_OpKernelContext* ctx, int i, const TF_Tensor* tensor,
+ TF_Status* status) {
+ auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
+ if (i < 0 || i >= cc_ctx->num_inputs()) {
+ TF_SetStatus(status, TF_OUT_OF_RANGE, "input index out of range");
+ return;
+ }
+ ::tensorflow::Tensor cc_tensor;
+ ::tensorflow::Status s = ::tensorflow::TF_TensorToTensor(tensor, &cc_tensor);
+ TF_SetStatus(status, TF_OK, "");
+ ::tensorflow::Set_TF_Status_from_Status(status, s);
+ if (s.ok()) {
+ cc_ctx->set_output(i, cc_tensor);
+ }
+}
diff --git a/tensorflow/c/kernels.h b/tensorflow/c/kernels.h
index d777882..1a91aa1 100644
--- a/tensorflow/c/kernels.h
+++ b/tensorflow/c/kernels.h
@@ -103,6 +103,14 @@
TF_CAPI_EXPORT extern void TF_GetInput(TF_OpKernelContext* ctx, int i,
TF_Tensor** tensor, TF_Status* status);
+// Sets the ith output of ctx to tensor. If TF_GetCode(status) is anything but
+// TF_OK, ctx is left unmodified.
+//
+// If i < 0 or i >= TF_NumOutputs(ctx), *status is set to TF_OUT_OF_RANGE.
+TF_CAPI_EXPORT extern void TF_SetOutput(TF_OpKernelContext* ctx, int i,
+ const TF_Tensor* tensor,
+ TF_Status* status);
+
#ifdef __cplusplus
} /* end extern "C" */
#endif
diff --git a/tensorflow/c/kernels_test.cc b/tensorflow/c/kernels_test.cc
index 80bf12c..e659ee3 100644
--- a/tensorflow/c/kernels_test.cc
+++ b/tensorflow/c/kernels_test.cc
@@ -147,6 +147,14 @@
EXPECT_EQ(TF_OUT_OF_RANGE, TF_GetCode(s));
TF_GetInput(ctx, 3, &input, s);
EXPECT_EQ(TF_OUT_OF_RANGE, TF_GetCode(s));
+
+ // Copy the input tensor to output.
+ TF_SetOutput(ctx, 0, input, s);
+ EXPECT_EQ(TF_OK, TF_GetCode(s));
+
+ TF_SetOutput(ctx, 24, input, s);
+ EXPECT_EQ(TF_OUT_OF_RANGE, TF_GetCode(s));
+
TF_DeleteStatus(s);
if (input != nullptr) {
TF_DeleteTensor(input);
@@ -188,6 +196,7 @@
ASSERT_EQ(2, num_inputs);
ASSERT_EQ(1, num_outputs);
+ ASSERT_EQ(123, ctx.mutable_output(0)->scalar<tensorflow::uint8>()());
}
}
diff --git a/tensorflow/compat_template_v1.__init__.py b/tensorflow/compat_template_v1.__init__.py
index 7df80ec..d58acde 100644
--- a/tensorflow/compat_template_v1.__init__.py
+++ b/tensorflow/compat_template_v1.__init__.py
@@ -23,12 +23,12 @@
# pylint: disable=g-bad-import-order
from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import
+# API IMPORTS PLACEHOLDER
+
from tensorflow.python.tools import component_api_helper as _component_api_helper
_component_api_helper.package_hook(
parent_package_str=__name__,
child_package_str=('tensorflow_estimator.python.estimator.api.estimator'))
-# API IMPORTS PLACEHOLDER
-
from tensorflow.python.platform import flags # pylint: disable=g-import-not-at-top
app.flags = flags # pylint: disable=undefined-variable
diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl
index 2dc3e8c..4051664 100644
--- a/tensorflow/compiler/aot/tfcompile.bzl
+++ b/tensorflow/compiler/aot/tfcompile.bzl
@@ -283,7 +283,7 @@
)
# Variables used for gen_test and gen_benchmark.
- cpp_class_split = cpp_class.rsplit("::", maxsplit = 2)
+ cpp_class_split = cpp_class.rsplit("::", 2)
if len(cpp_class_split) == 1:
no_ns_name = cpp_class_split[0]
else:
diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index be91ed4..d8c88a9 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -76,6 +76,7 @@
srcs = ["xla_cpu_device.cc"],
visibility = [":friends"],
deps = [
+ ":create_xla_launch_op", # buildcleaner: keep
":flags",
":jit_compilation_passes",
":xla_device",
@@ -95,6 +96,7 @@
srcs = ["xla_gpu_device.cc"],
visibility = [":friends"],
deps = [
+ ":create_xla_launch_op", # buildcleaner: keep
":jit_compilation_passes",
":xla_device",
"//tensorflow/compiler/jit/kernels:xla_ops",
@@ -104,6 +106,7 @@
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:lib",
"@com_google_absl//absl/memory",
+ "@com_google_absl//absl/strings",
],
alwayslink = 1,
)
@@ -512,6 +515,7 @@
"//tensorflow/compiler/jit/ops:xla_ops",
"//tensorflow/compiler/tf2xla:dump_graph",
"//tensorflow/compiler/tf2xla:resource_operation_table",
+ "//tensorflow/compiler/tf2xla:side_effect_util",
"//tensorflow/compiler/tf2xla:tf2xla_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla/cc:xla_jit_ops",
@@ -610,6 +614,7 @@
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:cc_ops_internal",
"//tensorflow/cc:function_ops",
+ "//tensorflow/cc:functional_ops",
"//tensorflow/cc:ops",
"//tensorflow/cc:resource_variable_ops",
"//tensorflow/cc:scope",
@@ -622,6 +627,7 @@
"//tensorflow/compiler/tf2xla/cc:xla_ops",
"//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
"//tensorflow/compiler/tf2xla/kernels:xla_ops",
+ "//tensorflow/compiler/xla:test",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc
index de89be9..8617beec 100644
--- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc
+++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc
@@ -299,7 +299,7 @@
.Attr("Toutputs: list(type) >= 0")
.Attr("ancestors: list(string) >= 0")
.Attr("key: string")
- .Attr("shape_inference_graph: string = ''")
+ .Attr("shape_inference_graph: func")
.Attr("shapes: list(shape) >= 0")
.SetShapeFn(::tensorflow::shape_inference::UnknownShape);
@@ -510,11 +510,7 @@
s = ConvertGraphDefToGraph(options, *graphdef, graph.get());
if (!s.ok()) return s;
- s = PerformStaticShapeInferenceBeforeEncapsulation(
- graph.get(), "_encapsulate", "_outside");
- if (!s.ok()) return s;
-
- s = PreprocessForEncapsulation(graph.get(), "_encapsulate", "_outside");
+ s = PerformStaticShapeInferenceBeforeEncapsulation(graph.get());
if (!s.ok()) return s;
std::unique_ptr<Graph> graph_out;
@@ -550,6 +546,14 @@
graphdef->Swap(&graphdef_out);
*library = lib_def->ToProto();
+ // Remove "_xla_inferred_shapes" attr. They are added by
+ // `PerformStaticShapeInferenceBeforeEncapsulation`.
+ for (FunctionDef& fdef : *library->mutable_function()) {
+ for (NodeDef& node_def : *fdef.mutable_node_def()) {
+ node_def.mutable_attr()->erase("_xla_inferred_shapes");
+ }
+ }
+
return s;
}
@@ -901,18 +905,22 @@
{
GraphDefBuilder shape(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape.opts());
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT, DT_FLOAT}, shape.opts());
+ Node* recv = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ shape.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1),
shape.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected));
}
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = test::function::XTimesTwo();
*library_expected.add_function() = FunctionDefHelper::Create(
"F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {},
@@ -931,8 +939,7 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph},
{"shapes", absl::Span<const DataType>({})},
{"_outside_compilation_subgraph", "O1"}},
{"c"}},
@@ -948,16 +955,18 @@
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT, DT_FLOAT}, b2.opts());
+ Node* recv = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1),
b2.opts()
.WithName("E")
- .WithControlInputs({recv, b})
+ .WithControlInputs({recv})
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
Node* send = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
- b2.opts().WithControlInput(e));
+ b2.opts().WithControlInput(e).WithAttr(
+ kXlaHasHostTransferAttrName, true));
Node* s = Sequencer(
b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}),
@@ -966,9 +975,9 @@
NodeBuilder node_builder("F1", "F1", lib_def.get());
node_builder.Input(a).Input(b);
Node* call =
- b2.opts().WithControlInputs({s}).FinalizeBuilder(&node_builder);
+ b2.opts().WithControlInputs({s, b}).FinalizeBuilder(&node_builder);
- Binary(a, call, b2.opts().WithName("G").WithControlInputs({e}));
+ Binary(a, call, b2.opts().WithName("G").WithControlInputs({call}));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -1022,14 +1031,16 @@
{
GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape1.opts());
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT, DT_FLOAT}, shape1.opts());
+ Node* recv = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1),
shape1.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
}
@@ -1037,33 +1048,45 @@
{
GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape2.opts());
- Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT, DT_FLOAT}, shape2.opts());
+ Node* recv1 = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1),
shape2.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2",
- {DT_FLOAT, DT_FLOAT}, shape2.opts());
+ Node* recv2 = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT, DT_FLOAT},
+ shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* g = Binary(e, ops::NodeOut(recv2, 0),
+ shape2.opts()
+ .WithName("G")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O2"));
Node* h = Binary(ops::NodeOut(recv2, 1), e,
shape2.opts()
.WithName("H")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O2"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {h}, shape2.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g, h},
+ shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape2, "F1_O2", &library_expected));
}
+ NameAttrList shape_inference_graph1, shape_inference_graph2;
+ shape_inference_graph1.set_name("_outside_compilation_shape_inference_F1_O1");
+ shape_inference_graph2.set_name("_outside_compilation_shape_inference_F1_O2");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"i_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"g_0_retval_retval:float", "i_0_retval_retval:float"}, {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}, {}},
{{"I"},
"UnaryTest",
- {"outside_compilation_O2_host_compute:outputs:0"}},
+ {"outside_compilation_O2_host_compute:outputs:1"}},
{{"F"},
"BinaryTest",
{"C:o:0", "outside_compilation_O1_host_compute:outputs:0"},
@@ -1073,11 +1096,10 @@
"XlaHostCompute",
{"F:o:0", "D:o:0"},
{{"Tinputs", absl::Span<const DataType>({DT_FLOAT, DT_FLOAT})},
- {"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"Toutputs", absl::Span<const DataType>({DT_FLOAT, DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O2"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O2"},
+ {"shape_inference_graph", shape_inference_graph2},
{"shapes", absl::Span<const DataType>({})},
{"_outside_compilation_subgraph", "O2"}},
{"F"}},
@@ -1088,13 +1110,13 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph1},
{"shapes", absl::Span<const DataType>({})},
{"_outside_compilation_subgraph", "O1"}},
{"D"}},
},
- {{"i_0_retval_retval", "I:o:0"}});
+ {{"g_0_retval_retval", "outside_compilation_O2_host_compute:outputs:0"},
+ {"i_0_retval_retval", "I:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -1105,19 +1127,22 @@
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT, DT_FLOAT}, b2.opts());
+ Node* recv1 = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1),
b2.opts()
.WithName("E")
- .WithControlInputs({recv1, b})
+ .WithControlInputs({recv1})
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
- b2.opts().WithControlInput(e));
+ b2.opts().WithControlInput(e).WithAttr(
+ kXlaHasHostTransferAttrName, true));
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2",
- {DT_FLOAT, DT_FLOAT}, b2.opts());
+ Node* recv2 = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT, DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* g = Binary(e, ops::NodeOut(recv2, 0),
b2.opts()
.WithName("G")
@@ -1130,7 +1155,8 @@
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O2"));
Node* send2 =
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {h}, b2.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g, h},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* s = Sequencer(b2.opts()
.WithName("F1_sequencer")
@@ -1139,12 +1165,13 @@
NodeBuilder node_builder("F1", "F1", lib_def.get());
node_builder.Input(a).Input(b);
- Node* call = b2.opts().WithControlInput(s).FinalizeBuilder(&node_builder);
+ Node* call =
+ b2.opts().WithControlInputs({s, b}).FinalizeBuilder(&node_builder);
- Binary(g, call, b2.opts().WithName("J"));
+ Binary(ops::NodeOut(call, 0), ops::NodeOut(call, 1),
+ b2.opts().WithName("J"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
-
TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef);
TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library);
}
@@ -1196,7 +1223,9 @@
*library_expected.add_function() = FunctionDefHelper::Create(
"F1", {"a_0_arg:float", "b_0_arg:float"},
- {"f_0_retval_retval:float", "d_0_retval_retval:float"}, {},
+ {"e_0_retval_retval:float", "f_0_retval_retval:float",
+ "d_0_retval_retval:float"},
+ {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
@@ -1212,35 +1241,37 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes",
absl::Span<const TensorShapeProto>({shape_proto_expected})},
{"_outside_compilation_subgraph", "O1"}},
{"D"}},
},
- {{"d_0_retval_retval", "D:o:0"}, {"f_0_retval_retval", "F:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"d_0_retval_retval", "D:o:0"},
+ {"f_0_retval_retval", "F:o:0"}});
*library_expected.add_function() = FunctionDefHelper::Create(
- "F2", {"f_0_arg:float", "bridge_e_g_0_arg:float"},
- {"i_0_retval_retval:float", "g_0_retval_retval:float"}, {},
+ "F2", {"e_0_arg:float", "f_0_arg:float", "d_0_arg:float"},
+ {"g_0_retval_retval:float", "i_0_retval_retval:float"}, {},
{
- {{"G"}, "BinaryTest", {"bridge_e_g_0_arg", "f_0_arg"}},
+ {{"G"}, "BinaryTest", {"e_0_arg", "f_0_arg"}},
{{"I"},
"BinaryTest",
{"f_0_arg", "outside_compilation_O1_host_compute:outputs:0"}},
{{"outside_compilation_O1_host_compute"},
"XlaHostCompute",
- {"G:o:0"},
- {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"d_0_arg", "G:o:0"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT, DT_FLOAT})},
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F2_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes",
absl::Span<const TensorShapeProto>({shape_proto_expected})},
{"_outside_compilation_subgraph", "O1"}}},
},
- {{"i_0_retval_retval", "I:o:0"}, {"g_0_retval_retval", "G:o:0"}});
+ {{"g_0_retval_retval", "G:o:0"}, {"i_0_retval_retval", "I:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -1251,16 +1282,18 @@
Node* key_constant1 =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv1 = RecvAtHost(ops::NodeOut(key_constant1, 0), "F1", "O1",
- {DT_FLOAT, DT_FLOAT}, b2.opts());
+ Node* recv1 = RecvAtHost(
+ ops::NodeOut(key_constant1, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1),
b2.opts()
.WithName("E")
- .WithControlInputs({recv1, b})
+ .WithControlInputs({recv1})
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
Node* send1 = SendFromHost(ops::NodeOut(key_constant1, 0), "F1", "O1", {e},
- b2.opts().WithControlInput(e));
+ b2.opts().WithControlInput(e).WithAttr(
+ kXlaHasHostTransferAttrName, true));
Node* s1 = Sequencer(
b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}),
"F1");
@@ -1268,29 +1301,33 @@
NodeBuilder node_builder1("F1", "F1", lib_def.get());
node_builder1.Input(a).Input(b);
Node* call1 =
- b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1);
+ b2.opts().WithControlInputs({s1, b}).FinalizeBuilder(&node_builder1);
Node* key_constant2 =
KeyPlaceholder("F2", b2.opts().WithName("F2_key_placeholder"));
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant2, 0), "F2", "O1",
- {DT_FLOAT}, b2.opts());
- Node* h = Binary(ops::NodeOut(call1, 1), recv2,
+ Node* recv2 = RecvAtHost(
+ ops::NodeOut(key_constant2, 0), "F2", "O1", {DT_FLOAT, DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* h = Binary(recv2, ops::NodeOut(recv2, 1),
b2.opts()
.WithName("H")
.WithAttr("_encapsulate", "F2")
.WithAttr("_outside", "O1"));
- Node* send2 = SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h},
- b2.opts());
+ Node* send2 =
+ SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* s2 = Sequencer(
b2.opts().WithName("F2_sequencer").WithControlInputs({recv2, send2}),
"F2");
NodeBuilder node_builder2("F2", "F2", lib_def.get());
- node_builder2.Input(call1).Input(e);
+ node_builder2.Input(call1)
+ .Input(ops::NodeOut(call1, 1))
+ .Input(ops::NodeOut(call1, 2));
Node* call2 = b2.opts()
- .WithControlInputs({s2, e, call1})
+ .WithControlInputs({s2, call1})
.FinalizeBuilder(&node_builder2);
- Binary(ops::NodeOut(call2, 1), call2, b2.opts().WithName("J"));
+ Binary(call2, ops::NodeOut(call2, 1), b2.opts().WithName("J"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -1326,8 +1363,7 @@
Node* h = Unary(g, b1.opts()
.WithName("H")
.WithAttr("_encapsulate", "F2")
- .WithAttr("_outside", "O1")
- .WithControlInput(e));
+ .WithAttr("_outside", "O1"));
Node* i = Unary(h, b1.opts().WithName("I").WithAttr("_encapsulate", "F2"));
Binary(f, i, b1.opts().WithName("J"));
TF_EXPECT_OK(b1.ToGraphDef(&graphdef));
@@ -1358,7 +1394,7 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes",
absl::Span<const TensorShapeProto>({shape_proto_expected})},
{"_outside_compilation_subgraph", "O1"}},
@@ -1380,7 +1416,7 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F2_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes",
absl::Span<const TensorShapeProto>({shape_proto_expected})},
{"_outside_compilation_subgraph", "O1"}}},
@@ -1401,7 +1437,7 @@
Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1),
b2.opts()
.WithName("E")
- .WithControlInputs({recv1, b})
+ .WithControlInputs({recv1})
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
Node* send1 = SendFromHost(ops::NodeOut(key_constant1, 0), "F1", "O1", {e},
@@ -1413,7 +1449,7 @@
NodeBuilder node_builder1("F1", "F1", lib_def.get());
node_builder1.Input(a).Input(b);
Node* call1 =
- b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1);
+ b2.opts().WithControlInputs({s1, b}).FinalizeBuilder(&node_builder1);
Node* key_constant2 =
KeyPlaceholder("F2", b2.opts().WithName("F2_key_placeholder"));
@@ -1422,8 +1458,7 @@
Node* h = Unary(recv2, b2.opts()
.WithName("H")
.WithAttr("_encapsulate", "F2")
- .WithAttr("_outside", "O1")
- .WithControlInput(e));
+ .WithAttr("_outside", "O1"));
Node* send2 = SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h},
b2.opts());
@@ -1484,12 +1519,12 @@
{"D:o:0", "outside_compilation_O1_host_compute:outputs:0"}},
{{"outside_compilation_O1_host_compute"},
"XlaHostCompute",
- {},
- {{"Tinputs", absl::Span<const DataType>({})},
+ {"a_0_arg"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes",
absl::Span<const TensorShapeProto>({shape_proto_expected})},
{"_outside_compilation_subgraph", "O1"}}},
@@ -1503,16 +1538,19 @@
Node* a = InputShaped(b2.opts().WithName("A"));
Node* b = Input(b2.opts().WithName("B"));
- Node* e = Unary(a, b2.opts()
- .WithName("E")
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O1"));
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
+ Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
+ {DT_FLOAT}, b2.opts());
+ Node* e = Unary(recv1, b2.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
Node* send1 =
SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts());
Node* s1 = Sequencer(
- b2.opts().WithName("F1_sequencer").WithControlInput(send1), "F1");
+ b2.opts().WithName("F1_sequencer").WithControlInputs({send1, recv1}),
+ "F1");
NodeBuilder node_builder1("F1", "F1", lib_def.get());
node_builder1.Input(a).Input(b);
Node* call1 =
@@ -1569,12 +1607,12 @@
{"D:o:0", "outside_compilation_O1_host_compute:outputs:0"}},
{{"outside_compilation_O1_host_compute"},
"XlaHostCompute",
- {},
- {{"Tinputs", absl::Span<const DataType>({})},
+ {"a_0_arg"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes",
absl::Span<const TensorShapeProto>({shape_proto_expected})},
{"_outside_compilation_subgraph", "O1"}},
@@ -1591,13 +1629,13 @@
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv1 =
- RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {}, b2.opts());
- Node* e = Unary(a, b2.opts()
- .WithName("E")
- .WithControlInput(recv1)
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O1"));
+ Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
+ {DT_FLOAT}, b2.opts());
+ Node* e = Unary(recv1, b2.opts()
+ .WithName("E")
+ .WithControlInput(recv1)
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
Node* send1 =
SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts());
Node* s1 = Sequencer(
@@ -1644,8 +1682,27 @@
FunctionDefLibrary library_expected;
GraphDef graphdef_expected;
+ {
+ GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
+ Node* key_constant = KeyPlaceholder("F1", shape1.opts());
+ Node* recv1 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = Unary(ops::NodeOut(recv1, 0), shape1.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ TF_EXPECT_OK(
+ AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
+ }
+
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"e_0_retval_retval:float", "f_0_retval_retval:float"}, {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
@@ -1654,14 +1711,15 @@
"XlaHostCompute",
{"D:o:0"},
{{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
- {"Toutputs", absl::Span<const DataType>({})},
+ {"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", shape_inference_graph},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O1"}}},
},
- {{"f_0_retval_retval", "F:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"f_0_retval_retval", "F:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -1678,14 +1736,17 @@
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
+ Node* send1 =
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts());
Node* s1 = Sequencer(
- b2.opts().WithName("F1_sequencer").WithControlInput(recv1), "F1");
+ b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}),
+ "F1");
NodeBuilder node_builder1("F1", "F1", lib_def.get());
node_builder1.Input(a).Input(b);
Node* call1 =
b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1);
- Binary(e, call1, b2.opts().WithName("G"));
+ Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("G"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -1722,8 +1783,27 @@
FunctionDefLibrary library_expected;
GraphDef graphdef_expected;
+ {
+ GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
+ Node* key_constant = KeyPlaceholder("F1", shape1.opts());
+ Node* recv1 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = Unary(ops::NodeOut(recv1, 0), shape1.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ TF_EXPECT_OK(
+ AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
+ }
+
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"e_0_retval_retval:float", "f_0_retval_retval:float"}, {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
@@ -1736,14 +1816,15 @@
"XlaHostCompute",
{"D:o:0"},
{{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
- {"Toutputs", absl::Span<const DataType>({})},
+ {"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", shape_inference_graph},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O1"}}},
},
- {{"f_0_retval_retval", "F:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"f_0_retval_retval", "F:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -1760,7 +1841,7 @@
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {},
+ Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
b2.opts().WithControlInput(e));
Node* s1 = Sequencer(
b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}),
@@ -1770,7 +1851,7 @@
Node* call1 =
b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1);
- Binary(e, call1, b2.opts().WithName("G"));
+ Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("G"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -1814,21 +1895,44 @@
GraphDef graphdef_expected;
{
+ GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
+ Node* key_constant = KeyPlaceholder("F1", shape1.opts());
+ Node* recv1 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = Unary(ops::NodeOut(recv1, 0), shape1.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ TF_EXPECT_OK(
+ AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
+ }
+
+ {
GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape2.opts());
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2",
- {DT_FLOAT}, shape2.opts());
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT},
+ shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* g = Unary(ops::NodeOut(recv2, 0), shape2.opts()
.WithName("G")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O2"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, shape2.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g},
+ shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape2, "F1_O2", &library_expected));
}
+ NameAttrList shape_inference_graph1;
+ shape_inference_graph1.set_name("_outside_compilation_shape_inference_F1_O1");
+ NameAttrList shape_inference_graph2;
+ shape_inference_graph2.set_name("_outside_compilation_shape_inference_F1_O2");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"e_0_retval_retval:float", "h_0_retval_retval:float"}, {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
@@ -1836,6 +1940,16 @@
{{"H"},
"UnaryTest",
{"outside_compilation_O2_host_compute:outputs:0"}},
+ {{"outside_compilation_O1_host_compute"},
+ "XlaHostCompute",
+ {"a_0_arg"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"ancestors", absl::Span<const string>({})},
+ {"key", "host_compute_channel_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph1},
+ {"shapes", absl::Span<const TensorShapeProto>({})},
+ {"_outside_compilation_subgraph", "O1"}}},
{{"outside_compilation_O2_host_compute"},
"XlaHostCompute",
{"F:o:0"},
@@ -1843,12 +1957,12 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O2"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O2"},
+ {"shape_inference_graph", shape_inference_graph2},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O2"}}},
},
- {{"h_0_retval_retval", "H:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"h_0_retval_retval", "H:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -1856,30 +1970,39 @@
GraphDefBuilder b2(GraphDefBuilder::kFailImmediately, lib_def.get());
Node* a = Input(b2.opts().WithName("A"));
Node* b = Input(b2.opts().WithName("B"));
-
- Node* e = Unary(a, b2.opts()
- .WithName("E")
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O1"));
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2",
- {DT_FLOAT}, b2.opts());
- Node* g = Unary(recv, b2.opts()
- .WithName("G")
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O2")
- .WithControlInput(e));
- Node* send =
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, b2.opts());
- Node* s1 = Sequencer(
- b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}),
- "F1");
+ Node* recv1 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+
+ Node* e = Unary(recv1, b2.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
+ Node* send1 =
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* g = Unary(recv2, b2.opts()
+ .WithName("G")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O2")
+ .WithControlInput(e));
+ Node* send2 =
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* s1 = Sequencer(b2.opts()
+ .WithName("F1_sequencer")
+ .WithControlInputs({recv1, send1, recv2, send2}),
+ "F1");
NodeBuilder node_builder1("F1", "F1", lib_def.get());
node_builder1.Input(a).Input(b).ControlInput(s1);
Node* call1 = b2.opts().FinalizeBuilder(&node_builder1);
- Binary(e, call1, b2.opts().WithName("I"));
+ Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("I"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -1925,19 +2048,24 @@
{
GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape1.opts());
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT}, shape1.opts());
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
}
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"e_0_retval_retval:float", "h_0_retval_retval:float"}, {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
@@ -1945,6 +2073,16 @@
"UnaryTest",
{"outside_compilation_O1_host_compute:outputs:0"}},
{{"H"}, "UnaryTest", {"F:o:0"}},
+ {{"outside_compilation_O2_host_compute"},
+ "XlaHostCompute",
+ {"a_0_arg"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"Toutputs", absl::Span<const DataType>({})},
+ {"ancestors", absl::Span<const string>({})},
+ {"key", "host_compute_channel_F1_O2"},
+ {"shape_inference_graph", NameAttrList()},
+ {"shapes", absl::Span<const TensorShapeProto>({})},
+ {"_outside_compilation_subgraph", "O2"}}},
{{"outside_compilation_O1_host_compute"},
"XlaHostCompute",
{"D:o:0"},
@@ -1952,12 +2090,12 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O1"}}},
},
- {{"h_0_retval_retval", "H:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"h_0_retval_retval", "H:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -1968,27 +2106,33 @@
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT}, b2.opts());
- Node* e = Unary(recv, b2.opts()
- .WithName("E")
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O1"));
+ Node* recv1 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = Unary(recv1, b2.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
Node* send =
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts());
- /*Node* g =*/Unary(a, b2.opts()
- .WithName("G")
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O2")
- .WithControlInput(e));
- Node* s1 = Sequencer(
- b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}),
- "F1");
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ /*Node* g =*/Unary(recv2, b2.opts()
+ .WithName("G")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O2")
+ .WithControlInput(e));
+ Node* s1 = Sequencer(b2.opts()
+ .WithName("F1_sequencer")
+ .WithControlInputs({recv1, recv2, send}),
+ "F1");
NodeBuilder node_builder1("F1", "F1", lib_def.get());
node_builder1.Input(a).Input(b).ControlInput(s1);
Node* call1 = b2.opts().FinalizeBuilder(&node_builder1);
- Binary(e, call1, b2.opts().WithName("I"));
+ Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("I"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -2039,19 +2183,24 @@
{
GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape1.opts());
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT}, shape1.opts());
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
}
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"e_0_retval_retval:float", "h_0_retval_retval:float"}, {},
{{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
{{"F"}, "UnaryTest", {"outside_compilation_O1_host_compute:outputs:0"}},
@@ -2063,8 +2212,7 @@
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O1"}}},
{{"outside_compilation_O2_host_compute"},
@@ -2074,7 +2222,7 @@
{"Toutputs", absl::Span<const DataType>({})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O2"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O2"}},
{}},
@@ -2085,11 +2233,12 @@
{"Toutputs", absl::Span<const DataType>({})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O3"},
- {"shape_inference_graph", ""},
+ {"shape_inference_graph", NameAttrList()},
{"shapes", absl::Span<const TensorShapeProto>({})},
{"_outside_compilation_subgraph", "O3"}},
{}}},
- {{"h_0_retval_retval", "H:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"h_0_retval_retval", "H:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -2100,23 +2249,27 @@
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT}, b2.opts());
+ Node* recv1 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* e = Unary(recv1, b2.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
Node* send =
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts());
- Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2",
- {DT_FLOAT}, b2.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
Node* g = Unary(recv2, b2.opts()
.WithName("G")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O2")
.WithControlInput(e));
- Node* recv3 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O3",
- {DT_FLOAT}, b2.opts());
+ Node* recv3 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O3", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
/*Node* i =*/Binary(recv3, e,
b2.opts()
.WithName("I")
@@ -2131,7 +2284,7 @@
node_builder1.Input(a).Input(b).ControlInput(s1);
Node* call1 = b2.opts().FinalizeBuilder(&node_builder1);
- Binary(e, call1, b2.opts().WithName("J"));
+ Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("J"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -2167,14 +2320,44 @@
FunctionDefLibrary library_expected;
GraphDef graphdef_expected;
+ {
+ GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately);
+ Node* key_constant = KeyPlaceholder("F1", shape1.opts());
+ Node* recv2 =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ TF_EXPECT_OK(
+ AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected));
+ }
+
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = FunctionDefHelper::Create(
- "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {},
+ "F1", {"a_0_arg:float", "b_0_arg:float"},
+ {"e_0_retval_retval:float", "f_0_retval_retval:float"}, {},
{
{{"C"}, "UnaryTest", {"a_0_arg"}},
{{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}},
{{"F"}, "UnaryTest", {"D:o:0"}},
+ {{"outside_compilation_O1_host_compute"},
+ "XlaHostCompute",
+ {"a_0_arg"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"ancestors", absl::Span<const string>({})},
+ {"key", "host_compute_channel_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph},
+ {"shapes", absl::Span<const TensorShapeProto>({})},
+ {"_outside_compilation_subgraph", "O1"}}},
},
- {{"f_0_retval_retval", "F:o:0"}});
+ {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"},
+ {"f_0_retval_retval", "F:o:0"}});
{
std::unique_ptr<FunctionLibraryDefinition> lib_def(
@@ -2183,15 +2366,26 @@
Node* a = Input(b2.opts().WithName("A"));
Node* b = Input(b2.opts().WithName("B"));
- Node* e = Unary(a, b2.opts()
- .WithName("E")
- .WithAttr("_encapsulate", "F1")
- .WithAttr("_outside", "O1"));
+ Node* key_constant =
+ KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
+ Node* recv =
+ RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = Unary(recv, b2.opts()
+ .WithName("E")
+ .WithAttr("_encapsulate", "F1")
+ .WithAttr("_outside", "O1"));
+ Node* send =
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* s = Sequencer(
+ b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}),
+ "F1");
NodeBuilder node_builder1("F1", "F1", lib_def.get());
- node_builder1.Input(a).Input(b);
+ node_builder1.Input(a).Input(b).ControlInput(s);
Node* call1 = b2.opts().FinalizeBuilder(&node_builder1);
- Binary(e, call1, b2.opts().WithName("G"));
+ Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("G"));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
@@ -2236,20 +2430,22 @@
{
GraphDefBuilder shape(GraphDefBuilder::kFailImmediately);
Node* key_constant = KeyPlaceholder("F1", shape.opts());
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT}, shape.opts());
- Node* a = InputShaped(shape.opts().WithName("A"));
- Node* c = Unary(a, shape.opts().WithName("C"));
- Node* e = BinaryUnknownShape(c, recv,
+ Node* recv = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ shape.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = BinaryUnknownShape(recv, ops::NodeOut(recv, 1),
shape.opts()
.WithName("E")
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
- SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts());
+ SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
+ shape.opts().WithAttr(kXlaHasHostTransferAttrName, true));
TF_EXPECT_OK(
AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected));
}
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1");
*library_expected.add_function() = test::function::XTimesTwo();
*library_expected.add_function() = FunctionDefHelper::Create(
"F1", {"b_0_arg:float", "c_0_arg:float"}, {"f_0_retval_retval:float"}, {},
@@ -2262,13 +2458,12 @@
{"outside_compilation_O1_host_compute"}},
{{"outside_compilation_O1_host_compute"},
"XlaHostCompute",
- {"c:o:0"},
- {{"Tinputs", absl::Span<const DataType>({DT_FLOAT})},
+ {"c_0_arg", "c:o:0"},
+ {{"Tinputs", absl::Span<const DataType>({DT_FLOAT, DT_FLOAT})},
{"Toutputs", absl::Span<const DataType>({DT_FLOAT})},
{"ancestors", absl::Span<const string>({})},
{"key", "host_compute_channel_F1_O1"},
- {"shape_inference_graph",
- "_outside_compilation_shape_inference_F1_O1"},
+ {"shape_inference_graph", shape_inference_graph},
{"shapes", absl::Span<const DataType>({})},
{"_outside_compilation_subgraph", "O1"}},
{"c"}},
@@ -2285,16 +2480,18 @@
Node* key_constant =
KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder"));
- Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1",
- {DT_FLOAT}, b2.opts());
- Node* e = BinaryUnknownShape(c, ops::NodeOut(recv, 0),
+ Node* recv = RecvAtHost(
+ ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT},
+ b2.opts().WithAttr(kXlaHasHostTransferAttrName, true));
+ Node* e = BinaryUnknownShape(recv, ops::NodeOut(recv, 1),
b2.opts()
.WithName("E")
- .WithControlInputs({recv, b})
+ .WithControlInputs({recv})
.WithAttr("_encapsulate", "F1")
.WithAttr("_outside", "O1"));
Node* send = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e},
- b2.opts().WithControlInput(e));
+ b2.opts().WithControlInput(e).WithAttr(
+ kXlaHasHostTransferAttrName, true));
Node* s = Sequencer(
b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}),
@@ -2303,9 +2500,9 @@
NodeBuilder node_builder("F1", "F1", lib_def.get());
node_builder.Input(b).Input(c);
Node* call =
- b2.opts().WithControlInputs({s, c}).FinalizeBuilder(&node_builder);
+ b2.opts().WithControlInputs({s, b, c}).FinalizeBuilder(&node_builder);
- Binary(a, call, b2.opts().WithName("G").WithControlInputs({e}));
+ Binary(a, call, b2.opts().WithName("G").WithControlInputs({call}));
TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected));
}
diff --git a/tensorflow/compiler/jit/encapsulate_util.cc b/tensorflow/compiler/jit/encapsulate_util.cc
index 1f4b9c9..2264806 100644
--- a/tensorflow/compiler/jit/encapsulate_util.cc
+++ b/tensorflow/compiler/jit/encapsulate_util.cc
@@ -62,517 +62,6 @@
n->AddAttr(attr_name, value);
}
-// Step 1a ~ 1d for PreprocessForEncapsulation(). See comments of
-// PreprocessForEncapsulation() for details.
-Status ProcessControlEdges(Graph* g, const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name) {
- // Gather edges to remove. We should not remove the edge while iterating.
- std::vector<const Edge*> edges_to_remove;
- for (const Edge* e : g->edges()) {
- if (!e->IsControlEdge()) {
- continue;
- }
-
- auto src_xla_computation =
- GetStringAttr(*e->src(), xla_computation_attr_name);
- auto dst_xla_computation =
- GetStringAttr(*e->dst(), xla_computation_attr_name);
- auto src_outside_compilation =
- GetStringAttr(*e->src(), outside_compilation_attr_name);
- auto dst_outside_compilation =
- GetStringAttr(*e->dst(), outside_compilation_attr_name);
-
- if (!src_xla_computation && !dst_xla_computation) {
- continue;
- } else if (src_xla_computation && !dst_xla_computation) {
- if (src_outside_compilation) {
- // Case 1c: outside compilation to host computation control edge.
- edges_to_remove.push_back(e);
-
- TF_RETURN_IF_ERROR(AppendToListAttr<string>(
- e->dst(), kXlaControlDependenciesAttrName, e->src()->name()));
- }
- } else if (!src_xla_computation && dst_xla_computation) {
- if (dst_outside_compilation) {
- // Case 1c: host computation control to outside compilation edge.
- edges_to_remove.push_back(e);
-
- TF_RETURN_IF_ERROR(AppendToListAttr<string>(
- e->dst(), kXlaControlDependenciesAttrName, e->src()->name()));
- }
- } else { // src_xla_computation && dst_xla_computation
- if (*src_xla_computation != *dst_xla_computation) {
- if (src_outside_compilation && dst_outside_compilation) {
- // Case 1b: outside compilation to outside compilation control edge.
- edges_to_remove.push_back(e);
-
- TF_RETURN_IF_ERROR(AppendToListAttr<string>(
- e->dst(), kXlaControlDependenciesAttrName, e->src()->name()));
- } else if (src_outside_compilation && !dst_outside_compilation) {
- // Case 1a: outside compilation to another XLA computaition control
- // edge.
- TF_RETURN_IF_ERROR(AppendToListAttr<string>(
- e->src(), kXlaConnectedToOtherXlaComputationAttrName,
- *dst_xla_computation));
- } else if (!src_outside_compilation && dst_outside_compilation) {
- // Case 1a: another XLA computaition to outside compilation control
- // edge.
- TF_RETURN_IF_ERROR(AppendToListAttr<string>(
- e->dst(), kXlaConnectedFromOtherXlaComputationAttrName,
- *src_xla_computation));
- }
- }
- }
- }
-
- for (auto e : edges_to_remove) {
- g->RemoveEdge(e);
- }
- return Status::OK();
-}
-
-// Step 2 for PreprocessForEncapsulation(). See comments of
-// PreprocessForEncapsulation() for details.
-Status ProcessXlaToXlaDataEdges(Graph* g,
- const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name) {
- // Gather edges between XLA computations. Notice that we do not store `Edge*`
- // directly because we remove some nodes while adding Identity nodes, and
- // those Edge pointers might be invalidated.
- struct EdgeInfo {
- int dst_input, dst_node_id;
- };
- std::vector<EdgeInfo> edges;
- for (const Edge* e : g->edges()) {
- if (e->IsControlEdge()) {
- continue;
- }
-
- auto src_xla_computation =
- GetStringAttr(*e->src(), xla_computation_attr_name);
- auto dst_xla_computation =
- GetStringAttr(*e->dst(), xla_computation_attr_name);
- auto src_outside_compilation =
- GetStringAttr(*e->src(), outside_compilation_attr_name);
- auto dst_outside_compilation =
- GetStringAttr(*e->dst(), outside_compilation_attr_name);
- if (!src_xla_computation || !dst_xla_computation) {
- continue;
- }
-
- if (*src_xla_computation != *dst_xla_computation) {
- if (src_outside_compilation || dst_outside_compilation) {
- edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id()});
- VLOG(4) << "XLA -> XLA edge: " << e->DebugString();
- }
- }
- }
-
- // For each XLA -> XLA edge, add an Identity node between src and dst.
- for (int i = 0; i < edges.size(); i++) {
- Node* dst = g->FindNodeId(edges[i].dst_node_id);
- const Edge* e;
- TF_RETURN_IF_ERROR(dst->input_edge(edges[i].dst_input, &e));
- Node* src = e->src();
- int src_output = e->src_output(), dst_input = e->dst_input();
- g->RemoveEdge(e);
-
- // Create Identity node, and connect it between `src` and `dst`.
- string identity_node_name =
- absl::StrCat("bridge_", src->name(), "_", dst->name());
- DataType dtype = src->output_type(src_output);
- TF_ASSIGN_OR_RETURN(Node * identity_node,
- BuildIdentityNode(g, identity_node_name, dtype, src,
- /*requested_device=*/absl::nullopt));
- identity_node->AddAttr(kBridgeSourceNodeAttrName, src->name());
- g->AddEdge(src, src_output, identity_node, 0);
- g->AddEdge(identity_node, 0, dst, dst_input);
-
- // Replace `e->dst()` because its input node changed.
- NodeDef new_def = dst->def();
- *new_def.mutable_input(dst_input) = identity_node->name();
- TF_ASSIGN_OR_RETURN(Node * dst_replace_node, ReplaceNode(g, dst, new_def));
-
- // Other edge in `edges` might have `e->dst()` as src or dst
- // node. Before removing `e->dst()`, replace those edges with corresponding
- // edges for `dst_replace_node`.
- for (int j = i + 1; j < edges.size(); j++) {
- if (edges[j].dst_node_id == edges[i].dst_node_id) {
- edges[j].dst_node_id = dst_replace_node->id();
- }
- }
- }
- return Status::OK();
-}
-
-// Step 3 for PreprocessForEncapsulation(). See comments of
-// PreprocessForEncapsulation() for details.
-Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation(
- Graph* g, const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name) {
- // Gather edges between outside compilation and host computation. Notice that
- // we do not store `Edge*` directly because we remove some nodes while adding
- // Identity nodes, and those Edge pointers might be invalidated.
- struct EdgeInfo {
- int dst_input, dst_node_id;
- bool is_host_to_outside_compilation;
- };
- std::vector<EdgeInfo> edges;
- for (const Edge* e : g->edges()) {
- if (e->IsControlEdge()) {
- continue;
- }
-
- if (e->src()->attrs().Find(xla_computation_attr_name) == nullptr &&
- e->dst()->attrs().Find(xla_computation_attr_name) != nullptr &&
- e->dst()->attrs().Find(outside_compilation_attr_name) != nullptr) {
- edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id(),
- /*is_host_to_outside_compilation=*/true});
- VLOG(4) << "Host -> oc edge: " << e->DebugString();
- } else if (e->dst()->attrs().Find(xla_computation_attr_name) == nullptr &&
- e->src()->attrs().Find(xla_computation_attr_name) != nullptr &&
- e->src()->attrs().Find(outside_compilation_attr_name) !=
- nullptr) {
- edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id(),
- /*is_host_to_outside_compilation=*/false});
- VLOG(4) << "Oc -> host edge: " << e->DebugString();
- }
- }
-
- // Remove the edge from host to outside compilation. Add a placeholder as
- // outside compilation node input.
- std::map<std::pair<string, int>, Node*> placeholders;
- for (int i = 0; i < edges.size(); i++) {
- Node* dst = g->FindNodeId(edges[i].dst_node_id);
- const Edge* e;
- TF_RETURN_IF_ERROR(dst->input_edge(edges[i].dst_input, &e));
- Node* src = e->src();
- int src_output = e->src_output(), dst_input = e->dst_input();
- g->RemoveEdge(e);
-
- // Find or create placeholder node.
- string new_name =
- edges[i].is_host_to_outside_compilation
- ? absl::StrCat(src->name(), "_host_to_oc_placeholder_", src_output)
- : absl::StrCat(src->name(), "_oc_to_host_placeholder_", src_output);
- auto placeholder_index = std::make_pair(src->name(), src_output);
- auto iter = placeholders.find(placeholder_index);
- Node* placeholder_node;
- if (iter == placeholders.end()) {
- NodeDefBuilder placeholder_builder(new_name, "Placeholder");
- placeholder_builder.Attr("dtype", src->output_type(src_output));
- if (edges[i].is_host_to_outside_compilation) {
- placeholder_builder.Attr(kHostToOutsideCompilationOriginalNodeAttrName,
- src->name());
- placeholder_builder.Attr(kHostToOutsideCompilationSrcOutputAttrName,
- src_output);
- // If this placeholder node is in outside compilation, we need to set
- // `xla_computation_attr_name` and `outside_compilation_attr_name`.
- string xla_computation_attr, outside_compilation_attr;
- TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(), xla_computation_attr_name,
- &xla_computation_attr));
- TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(),
- outside_compilation_attr_name,
- &outside_compilation_attr));
- placeholder_builder.Attr(xla_computation_attr_name,
- xla_computation_attr);
- placeholder_builder.Attr(outside_compilation_attr_name,
- outside_compilation_attr);
- } else {
- placeholder_builder.Attr(kOutsideCompilationToHostOriginalNodeAttrName,
- src->name());
- placeholder_builder.Attr(kOutsideCompilationToHostSrcOutputAttrName,
- src_output);
- }
- NodeDef placeholder_def;
- TF_RETURN_IF_ERROR(placeholder_builder.Finalize(&placeholder_def));
- Status s;
- placeholder_node = g->AddNode(placeholder_def, &s);
- TF_RETURN_IF_ERROR(s);
- placeholders[placeholder_index] = placeholder_node;
- } else {
- placeholder_node = iter->second;
- }
- g->AddEdge(placeholder_node, 0, dst, dst_input);
-
- // Replace `e->dst()` because its input node changed.
- NodeDef new_def = dst->def();
- *new_def.mutable_input(dst_input) = placeholder_node->name();
- TF_ASSIGN_OR_RETURN(Node * dst_replace_node, ReplaceNode(g, dst, new_def));
-
- // Other edge in `edges` might have `e->dst()` as src or dst
- // node. Before removing `e->dst()`, replace those edges with corresponding
- // edges for `dst_replace_node`.
- for (int j = i + 1; j < edges.size(); j++) {
- if (edges[j].dst_node_id == edges[i].dst_node_id) {
- edges[j].dst_node_id = dst_replace_node->id();
- }
- }
- }
- return Status::OK();
-}
-
-// Step 1 for `PostprocessForEncapsulation`. See comments of
-// `PostprocessForEncapsulation` for details.
-Status RemovePlaceholderBetweenOutsideCompilationAndHostComputation(Graph* g) {
- // Gather all outside compilation to host computation nodes.
- struct PlaceHolderNodeInfo {
- Node* n;
- bool is_host_to_oc;
- };
- std::vector<PlaceHolderNodeInfo> placeholder_nodes;
- for (Node* n : g->nodes()) {
- if (n->type_string() == "Placeholder") {
- if (HasNodeAttr(n->def(),
- kOutsideCompilationToHostOriginalNodeAttrName)) {
- placeholder_nodes.push_back({n, false});
- } else if (HasNodeAttr(n->def(),
- kHostToOutsideCompilationOriginalNodeAttrName)) {
- placeholder_nodes.push_back({n, true});
- }
- }
- }
-
- // Remove the placeholder nodes, and reconnect original edge.
- auto node_name_index = g->BuildNodeNameIndex();
- for (auto placeholder_iter : placeholder_nodes) {
- Node* n = placeholder_iter.n;
-
- string node_name;
- int node_src_output;
- if (placeholder_iter.is_host_to_oc) {
- TF_RETURN_IF_ERROR(
- GetNodeAttr(n->attrs(), kHostToOutsideCompilationOriginalNodeAttrName,
- &node_name));
- TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(),
- kHostToOutsideCompilationSrcOutputAttrName,
- &node_src_output));
- } else {
- TF_RETURN_IF_ERROR(
- GetNodeAttr(n->attrs(), kOutsideCompilationToHostOriginalNodeAttrName,
- &node_name));
- TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(),
- kOutsideCompilationToHostSrcOutputAttrName,
- &node_src_output));
- }
- auto iter = node_name_index.find(node_name);
- if (iter == node_name_index.end()) {
- return errors::Internal(
- "Cannot find original node for oc -> host placeholder node ",
- node_name);
- }
-
- // Change all usage node to use the original node instead.
- Node* original_node = iter->second;
- std::vector<const Edge*> control_edges;
- std::vector<OutEdgeInfo> data_edges;
- for (auto e : n->out_edges()) {
- if (e->IsControlEdge()) {
- control_edges.push_back(e);
- } else {
- data_edges.push_back({e->dst(), e->src_output(), e->dst_input()});
- }
- }
- for (const Edge* e : control_edges) {
- g->AddControlEdge(original_node, e->dst());
- g->RemoveEdge(e);
- }
- for (int i = 0; i < data_edges.size(); i++) {
- Node* dst = data_edges[i].dst;
- NodeDef new_def = dst->def();
- int dst_input = data_edges[i].dst_input;
- *new_def.mutable_input(dst_input) =
- absl::StrCat(original_node->name(), ":", node_src_output);
- TF_ASSIGN_OR_RETURN(Node * replace_node, ReplaceNode(g, dst, new_def));
-
- const Edge* edge_to_replace = nullptr;
- TF_RETURN_IF_ERROR(replace_node->input_edge(dst_input, &edge_to_replace));
- g->RemoveEdge(edge_to_replace);
- g->AddEdge(original_node, node_src_output, replace_node, dst_input);
-
- // Other edges might have `dst` as dst node. Update those edges with
- // `replace_node`.
- for (int j = i + 1; j < data_edges.size(); j++) {
- if (data_edges[j].dst == dst) {
- data_edges[j].dst = replace_node;
- }
- }
-
- // Other placeholder node might have `dst` as original node. Update
- // `node_name_index` with `replace_node`.
- node_name_index[replace_node->name()] = replace_node;
- }
-
- // Remove placeholder node.
- g->RemoveNode(n);
- }
- return Status::OK();
-}
-
-// Step 2 for `PostprocessForEncapsulation`. See comments of
-// `PostprocessForEncapsulation` for details.
-Status RemoveIdentityBetweenDifferentXlaComputation(Graph* g) {
- // Gather Identity nodes to remove.
- std::vector<Node*> bridge_nodes;
- for (Node* n : g->nodes()) {
- if (n->type_string() == "Identity" &&
- HasNodeAttr(n->def(), kBridgeSourceNodeAttrName)) {
- bridge_nodes.push_back(n);
- }
- }
-
- // Remove the identity nodes, and reconnect the original edge.
- for (int i = 0; i < bridge_nodes.size(); i++) {
- Node* n = bridge_nodes[i];
- const Edge* src_edge = nullptr;
- TF_RETURN_IF_ERROR(n->input_edge(0, &src_edge));
-
- // Change all usage node to use the original node instead.
- std::vector<const Edge*> control_edges;
- std::vector<OutEdgeInfo> data_edges;
- for (auto e : n->out_edges()) {
- if (e->IsControlEdge()) {
- control_edges.push_back(e);
- } else {
- data_edges.push_back({e->dst(), e->src_output(), e->dst_input()});
- }
- }
- for (const Edge* e : control_edges) {
- g->AddControlEdge(src_edge->src(), e->dst());
- g->RemoveEdge(e);
- }
- for (int j = 0; j < data_edges.size(); j++) {
- Node* dst = data_edges[j].dst;
- NodeDef new_def = dst->def();
- int dst_input = data_edges[j].dst_input;
- *new_def.mutable_input(dst_input) =
- absl::StrCat(src_edge->src()->name(), ":", src_edge->src_output());
- TF_ASSIGN_OR_RETURN(Node * replace_node, ReplaceNode(g, dst, new_def));
-
- const Edge* edge_to_replace = nullptr;
- TF_RETURN_IF_ERROR(replace_node->input_edge(dst_input, &edge_to_replace));
- g->RemoveEdge(edge_to_replace);
- g->AddEdge(src_edge->src(), src_edge->src_output(), replace_node,
- dst_input);
-
- // Other edges might have `dst` as dst node. Update those edges with
- // `replace_node`.
- for (int k = j + 1; k < data_edges.size(); k++) {
- if (data_edges[k].dst == dst) {
- data_edges[k].dst = replace_node;
- }
- }
-
- // The node we replaced might be in `bridge_nodes`. If so, update
- // `bridge_nodes` to use the replaced node.
- for (int k = i + 1; k < bridge_nodes.size(); k++) {
- if (bridge_nodes[k] == dst) {
- bridge_nodes[k] = replace_node;
- }
- }
- }
-
- // Remove Identity node.
- g->RemoveNode(n);
- }
- return Status::OK();
-}
-
-// Step 3 for `PostprocessForEncapsulation`. See comments of
-// `PostprocessForEncapsulation` for details.
-// We do not need to worry about removed nodes in step 1 and 2;
-// `PreprocessForEncapsulation` will not record control dependencies for those
-// remvoed nodes in the first place.
-Status AddControlDependencies(
- Graph* g, const std::unordered_map<string, string>& cluster_node_names) {
- auto node_name_index = g->BuildNodeNameIndex();
-
- // Reconnect outside compilation to outside compilation control edge.
- for (Node* n : g->nodes()) {
- std::vector<string> control_deps;
- Status s =
- GetNodeAttr(n->attrs(), kXlaControlDependenciesAttrName, &control_deps);
- if (!s.ok()) {
- if (s.code() != error::NOT_FOUND) {
- return s;
- } else {
- continue;
- }
- } else {
- n->ClearAttr(kXlaControlDependenciesAttrName);
- for (const string& control_input : control_deps) {
- auto iter = node_name_index.find(control_input);
- if (iter == node_name_index.end()) {
- return errors::Internal("Cannot find original node for ",
- control_input);
- }
- g->AddControlEdge(iter->second, n);
- }
- }
- }
-
- // Reconnect outside compilation to XLA computation control edge.
- for (Node* n : g->nodes()) {
- std::vector<string> control_deps;
- Status s = GetNodeAttr(
- n->attrs(), kXlaConnectedToOtherXlaComputationAttrName, &control_deps);
- if (!s.ok()) {
- if (s.code() != error::NOT_FOUND) {
- return s;
- } else {
- continue;
- }
- } else {
- n->ClearAttr(kXlaConnectedToOtherXlaComputationAttrName);
- for (const string& control_input : control_deps) {
- auto iter = cluster_node_names.find(control_input);
- if (iter == cluster_node_names.end()) {
- return errors::Internal("Cannot find cluster node for ",
- control_input);
- }
- auto iter2 = node_name_index.find(iter->second);
- if (iter2 == node_name_index.end()) {
- return errors::Internal("Cannot find cluster node for ",
- iter->second);
- }
- g->AddControlEdge(n, iter2->second);
- }
- }
- }
-
- // Reconnect XLA computation to outside compilation control edge.
- for (Node* n : g->nodes()) {
- std::vector<string> control_deps;
- Status s =
- GetNodeAttr(n->attrs(), kXlaConnectedFromOtherXlaComputationAttrName,
- &control_deps);
- if (!s.ok()) {
- if (s.code() != error::NOT_FOUND) {
- return s;
- } else {
- continue;
- }
- } else {
- n->ClearAttr(kXlaConnectedFromOtherXlaComputationAttrName);
- for (const string& control_input : control_deps) {
- auto iter = cluster_node_names.find(control_input);
- if (iter == cluster_node_names.end()) {
- return errors::Internal("Cannot find cluster node for ",
- control_input);
- }
- auto iter2 = node_name_index.find(iter->second);
- if (iter2 == node_name_index.end()) {
- return errors::Internal("Cannot find cluster node for ",
- iter->second);
- }
- g->AddControlEdge(iter2->second, n);
- }
- }
- }
-
- return Status::OK();
-}
-
// Step 1 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of
// `PreprocessEdgesBetweenOutsideCompilations` for details.
Status PreprocessControlEdgesBetweenOutsideCompilations(
@@ -811,20 +300,6 @@
const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes";
-const char kXlaConnectedToOtherXlaComputationAttrName[] =
- "_xla_connected_to_other_xla_computation";
-const char kXlaConnectedFromOtherXlaComputationAttrName[] =
- "_xla_connected_from_other_xla_computation";
-const char kXlaControlDependenciesAttrName[] = "_xla_control_dependencies";
-const char kBridgeSourceNodeAttrName[] = "_xla_bridge_src";
-const char kOutsideCompilationToHostOriginalNodeAttrName[] =
- "_xla_oc_to_host_node_name";
-const char kOutsideCompilationToHostSrcOutputAttrName[] =
- "_xla_oc_to_host_src_output";
-const char kHostToOutsideCompilationOriginalNodeAttrName[] =
- "_xla_host_to_oc_node_name";
-const char kHostToOutsideCompilationSrcOutputAttrName[] =
- "_xla_host_to_oc_src_output";
const char kXlaConnectedToXlaComputationAttrName[] =
"_xla_connected_to_xla_computation";
const char kXlaConnectedFromXlaComputationAttrName[] =
@@ -835,32 +310,7 @@
const char kXlaControlDependenciesWithinXlaClusterAttrName[] =
"_xla_control_dependencies_within_xla_cluster";
-Status PerformStaticShapeInferenceBeforeEncapsulation(
- Graph* g, const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name) {
- // Find all outside compilation to XLA computation data edges.
- std::unordered_set<Node*> outside_compilation_send_nodes;
- for (auto e : g->edges()) {
- if (e->IsControlEdge()) {
- continue;
- }
-
- auto src_computation = GetStringAttr(*e->src(), xla_computation_attr_name);
- auto dst_computation = GetStringAttr(*e->dst(), xla_computation_attr_name);
- if (!src_computation || !dst_computation ||
- *src_computation != *dst_computation) {
- continue;
- }
-
- auto src_outside_compilation =
- GetStringAttr(*e->src(), outside_compilation_attr_name);
- auto dst_outside_compilation =
- GetStringAttr(*e->dst(), outside_compilation_attr_name);
- if (src_outside_compilation && !dst_outside_compilation) {
- outside_compilation_send_nodes.insert(e->src());
- }
- }
-
+Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g) {
// Perform shape inference.
std::map<int, InferredShape> arg_shapes;
GraphShapeInfo shape_info;
@@ -868,55 +318,21 @@
InferShapes(g, arg_shapes, /*fnlib_def=*/nullptr, &shape_info));
// Add attribute for output shapes.
- for (Node* n : outside_compilation_send_nodes) {
- auto iter = shape_info.find(n->name());
- if (iter == shape_info.end()) {
- continue;
- }
-
+ auto node_name_index = g->BuildNodeNameIndex();
+ for (auto iter : shape_info) {
std::vector<PartialTensorShape> output_shapes;
- std::transform(iter->second.begin(), iter->second.end(),
+ std::transform(iter.second.begin(), iter.second.end(),
std::back_inserter(output_shapes),
[](const InferredShape& inferred_shape) {
return inferred_shape.shape;
});
+ Node* n = node_name_index[iter.first];
n->AddAttr(kXlaInferredShapesAttrName, output_shapes);
}
return Status::OK();
}
-Status PreprocessForEncapsulation(Graph* g,
- const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name) {
- TF_RETURN_IF_ERROR(ProcessControlEdges(g, xla_computation_attr_name,
- outside_compilation_attr_name));
- TF_RETURN_IF_ERROR(ProcessXlaToXlaDataEdges(g, xla_computation_attr_name,
- outside_compilation_attr_name));
- TF_RETURN_IF_ERROR(ProcessDataEdgeBetweenOutsideCompilationAndHostComputation(
- g, xla_computation_attr_name, outside_compilation_attr_name));
- return Status::OK();
-}
-
-Status PostprocessForEncapsulation(
- Graph* g, const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name,
- const std::unordered_map<string, XlaClusterInfo>& clusters) {
- // The `node` pointer in `XlaClusterInfo` might be invalidated in step 1/2,
- // but the node name won't change. Record cluster node name for
- // `AddControlDependencies`.
- std::unordered_map<string, string> cluster_node_names;
- for (const auto& iter : clusters) {
- cluster_node_names[iter.first] = iter.second.node->name();
- }
-
- TF_RETURN_IF_ERROR(
- RemovePlaceholderBetweenOutsideCompilationAndHostComputation(g));
- TF_RETURN_IF_ERROR(RemoveIdentityBetweenDifferentXlaComputation(g));
- TF_RETURN_IF_ERROR(AddControlDependencies(g, cluster_node_names));
- return Status::OK();
-}
-
Status PreprocessEdgesBetweenOutsideCompilations(
Graph* g, const string& outside_compilation_attr_name) {
// Remove edges from source node to outside compilation nodes, and edges
diff --git a/tensorflow/compiler/jit/encapsulate_util.h b/tensorflow/compiler/jit/encapsulate_util.h
index e363bc5..c9f16d1 100644
--- a/tensorflow/compiler/jit/encapsulate_util.h
+++ b/tensorflow/compiler/jit/encapsulate_util.h
@@ -27,51 +27,13 @@
// a list of PartialTensorShape objects.
extern const char kXlaInferredShapesAttrName[];
-// Infer output shapes for outside compilation nodes which have output data
-// edges to XLA computation nodes. These shapes will be used later by XLA
-// compiler as output shapes of the outside compilation's XlaHostCompute op.
-// XLA computation nodes will be mark by attr `xla_computation_attr_name`;
-// outside compilation nodes will be marked by both attr
-// `xla_computation_attr_name` and `outside_compilation_attr_name`.
-//
-// Those outside compilation nodes will be marked with attribute
-// `kXlaInferredShapesAttrName`.
+// Infers output shapes for all nodes in graph `g`. The output shapes will be
+// stored in node attribute `kXlaInferredShapesAttrName`.
//
// We have to perform shape inference before encapsulation because after
// encapsulation, some nodes will be encapsulated into function call, and shape
// inference does not handle function call at the moment.
-Status PerformStaticShapeInferenceBeforeEncapsulation(
- Graph* g, const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name);
-
-// Attribute indicating that some ops in other XLA computation has control
-// dependency on this node. Attribute value will be a list of string (XLA
-// computation names).
-extern const char kXlaConnectedToOtherXlaComputationAttrName[];
-
-// Attribute indicating that this node has control dependency on some ops in
-// other XLA computation. Attribute value will be a list of string (XLA
-// computation names).
-extern const char kXlaConnectedFromOtherXlaComputationAttrName[];
-
-// Attribute indicating that this node has control dependencies on some other
-// nodes. Attribute value will be a list of string (node names).
-extern const char kXlaControlDependenciesAttrName[];
-
-// Attribute indicating that this is an Identity node added to act as a bridge
-// between different XLA computations. Attribute value will be string (source
-// node name).
-extern const char kBridgeSourceNodeAttrName[];
-
-// Attribute indicating that this is an Placeholder node added to act as a
-// temporary input node for an outside compilation node. Attribute value will be
-// string (original input node name).
-extern const char kOutsideCompilationToHostOriginalNodeAttrName[];
-
-// Attribute indicating that this is an Placeholder node added to act as a
-// temporary input node for an outside compilation node. Attribute value will be
-// int (src_output for original edge).
-extern const char kOutsideCompilationToHostSrcOutputAttrName[];
+Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g);
// Attribute indicating that some ops in this node's XLA computation has control
// dependency on this node. Attribute value will always be "true".
@@ -82,16 +44,6 @@
extern const char kXlaConnectedFromXlaComputationAttrName[];
// Attribute indicating that this is an Placeholder node added to act as a
-// temporary input node for an host node. Attribute value will be string
-// (original input node name).
-extern const char kHostToOutsideCompilationOriginalNodeAttrName[];
-
-// Attribute indicating that this is an Placeholder node added to act as a
-// temporary input node for a host node. Attribute value will be int (src_output
-// for original edge).
-extern const char kHostToOutsideCompilationSrcOutputAttrName[];
-
-// Attribute indicating that this is an Placeholder node added to act as a
// temporary input node for an outside compilation node. Attribute value will be
// string (original input node name).
extern const char kOutsideCompilationOriginalNodeAttrName[];
@@ -106,27 +58,6 @@
// (node names).
extern const char kXlaControlDependenciesWithinXlaClusterAttrName[];
-// Preprocesses edges between different XLA clusters for encapsulation. It will
-// perform the following operations in order:
-//
-// 1a. For control edges between outside compilation and another XLA
-// computation, add attr "kXlaConnected{From, To}OtherXlaComputationAttrName
-// = XLA computation node name" to the outside compilation node.
-// 1b. For control edges between different outside compilations (in different
-// XLA computations), remove the edge and add attr
-// "kXlaControlDependenciesAttrName = src node name" to dst node.
-// 1c. For control edges between outside compilation and host computation,
-// remove the edge and add attr "kXlaControlDependenciesAttrName = src node
-// name" to dst node.
-// 2. For data edges between different XLA computations, if either src or dst
-// is outside compilation, add an Identity node in between the edge. The
-// identity node will have attr kBridgeSourceNodeAttrName.
-// 3. For data edges between outside compilation and host computation, remove
-// the edge and create a Placeholder node as dst node's input.
-Status PreprocessForEncapsulation(Graph* g,
- const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name);
-
// Information for XLA computation.
struct XlaClusterInfo {
// Add an explicitly-defined default constructor for this class.
@@ -158,24 +89,6 @@
const std::map<string, int> host_compute_core;
};
-// Postprocesses edges between different XLA clusters for encapsulation. This
-// function reverts what `PreprocessForEncapsulation` did. It will perform the
-// following operations in order:
-//
-// 1. Remove Placeholder nodes between outside compilation and host computation
-// (created in `PreprocessForEncapsulation` step 3).
-// 2. Remove Identity nodes created in `PreprocessForEncapsulation` step 2.
-// 3a. Reconnect control edges between outside compilation and another XLA
-// computation (marked by `PreprocessForEncapsulation` step 1a).
-// 3b. Reconnect control edges between different outside compilations (marked by
-// `PreprocessForEncapsulation` step 1b).
-// 3c. Reconnect control edges between outside compilation and host computation
-// (marked by `PreprocessForEncapsulation` step 1c).
-Status PostprocessForEncapsulation(
- Graph* g, const string& xla_computation_attr_name,
- const string& outside_compilation_attr_name,
- const std::unordered_map<string, XlaClusterInfo>& clusters);
-
// Preprocesses edges within the same XLA cluster. It will perform the following
// operations in order:
//
diff --git a/tensorflow/compiler/jit/encapsulate_util_test.cc b/tensorflow/compiler/jit/encapsulate_util_test.cc
index 3b8b49c..3bb979e 100644
--- a/tensorflow/compiler/jit/encapsulate_util_test.cc
+++ b/tensorflow/compiler/jit/encapsulate_util_test.cc
@@ -38,24 +38,11 @@
Graph g(OpRegistry::Global());
TF_CHECK_OK(s.ToGraph(&g));
- // "add" node is outside compilation node, "identity" node is XLA node.
- auto node_index = g.BuildNodeNameIndex();
- Node *add_node = node_index["add"], *identity_node = node_index["identity"];
- add_node->AddAttr("_xla", "cluster");
- add_node->AddAttr("_oc", "cluster");
- identity_node->AddAttr("_xla", "cluster");
- TF_CHECK_OK(
- PerformStaticShapeInferenceBeforeEncapsulation(&g, "_xla", "_oc"));
+ TF_CHECK_OK(PerformStaticShapeInferenceBeforeEncapsulation(&g));
- // Check that only "add" node now has _xla_inferred_shapes attr.
- std::vector<Node *> nodes_with_inferred_shape;
- for (Node *n : g.nodes()) {
- if (HasNodeAttr(n->def(), kXlaInferredShapesAttrName)) {
- nodes_with_inferred_shape.push_back(n);
- }
- }
- EXPECT_EQ(nodes_with_inferred_shape.size(), 1);
- EXPECT_EQ(nodes_with_inferred_shape[0], add_node);
+ // Check that "add" node now has _xla_inferred_shapes attr.
+ auto node_index = g.BuildNodeNameIndex();
+ Node *add_node = node_index["add"];
std::vector<PartialTensorShape> output_shapes;
TF_CHECK_OK(GetNodeAttr(add_node->attrs(), kXlaInferredShapesAttrName,
&output_shapes));
@@ -66,329 +53,4 @@
EXPECT_EQ(shape_proto.dim(0).size(), 2);
}
-TEST(PreprocessForEncapsulationTest, ControlEdges) {
- // Build the graph:
- // "const_0" and "const_1" in host computation
- // "add" = "const_0" + "const_1" in XLA computation 0
- // "identity0" = "add" in XLA computation 0 & outside compilation 0
- // "identity1" = "identity0" in XLA computation 0
- // "identity2" = "identity1" in host computation
- // "identity3" = "identity2" in XLA computation 1
- // "identity4" = "identity3" in XLA computation 1 & outside compilation 1
- // "identity5" = "identity4" in XLA computation 1
- // "identity6" = "identity5" in host computation
- tensorflow::Scope s = tensorflow::Scope::NewRootScope();
- Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {});
- Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {});
- Output add = ops::Add(s.WithOpName("add"), const_0, const_1);
- Output identity0 = ops::Identity(s.WithOpName("identity0"), add);
- Output identity1 = ops::Identity(s.WithOpName("identity1"), identity0);
- Output identity2 = ops::Identity(s.WithOpName("identity2"), identity1);
- Output identity3 = ops::Identity(s.WithOpName("identity3"), identity2);
- Output identity4 = ops::Identity(s.WithOpName("identity4"), identity3);
- Output identity5 = ops::Identity(s.WithOpName("identity5"), identity4);
- Graph g(OpRegistry::Global());
- TF_CHECK_OK(s.ToGraph(&g));
- auto node_index = g.BuildNodeNameIndex();
-
- // Set XLA computation/outside compilation attr, and add control edges.
- Node *const0_node = node_index["const_0"], *add_node = node_index["add"],
- *identity0_node = node_index["identity0"],
- *identity1_node = node_index["identity1"],
- *identity2_node = node_index["identity2"],
- *identity3_node = node_index["identity3"],
- *identity4_node = node_index["identity4"],
- *identity5_node = node_index["identity5"];
- add_node->AddAttr("_xla", "0");
- identity0_node->AddAttr("_xla", "0");
- identity0_node->AddAttr("_oc", "0");
- identity1_node->AddAttr("_xla", "0");
- identity3_node->AddAttr("_xla", "1");
- identity4_node->AddAttr("_xla", "1");
- identity4_node->AddAttr("_oc", "0");
- identity5_node->AddAttr("_xla", "1");
- // Case 1a: control edges between outside compilation and another XLA
- // computation.
- g.AddControlEdge(identity0_node, identity3_node);
- g.AddControlEdge(identity1_node, identity4_node);
- // Case 1b: control edges between different outside compilations.
- g.AddControlEdge(identity0_node, identity4_node);
- // Case 1c: control edges between outside compilation and host computation.
- g.AddControlEdge(const0_node, identity0_node);
- g.AddControlEdge(identity0_node, identity2_node);
-
- TF_CHECK_OK(PreprocessForEncapsulation(&g, "_xla", "_oc"));
-
- // Case 1a: add attr "_xla_control_deps_{from/to} = XLA computation node name"
- // to the outside compilation node.
- std::vector<string> attr;
- TF_CHECK_OK(GetNodeAttr(identity0_node->def(),
- kXlaConnectedToOtherXlaComputationAttrName, &attr));
- EXPECT_EQ(attr.size(), 1);
- EXPECT_EQ(attr[0], "1");
- attr.clear();
- TF_CHECK_OK(GetNodeAttr(identity4_node->def(),
- kXlaConnectedFromOtherXlaComputationAttrName, &attr));
- EXPECT_EQ(attr.size(), 1);
- EXPECT_EQ(attr[0], "0");
- // Case 1b: add attr "_xla_control_deps = src node name" to dst node.
- attr.clear();
- TF_CHECK_OK(GetNodeAttr(identity4_node->def(),
- kXlaControlDependenciesAttrName, &attr));
- EXPECT_EQ(attr.size(), 1);
- EXPECT_EQ(attr[0], "identity0");
- // Case 1c: add attr "_xla_control_deps = src node name" to dst node.
- attr.clear();
- TF_CHECK_OK(GetNodeAttr(identity0_node->def(),
- kXlaControlDependenciesAttrName, &attr));
- EXPECT_EQ(attr.size(), 1);
- EXPECT_EQ(attr[0], "const_0");
- attr.clear();
- TF_CHECK_OK(GetNodeAttr(identity2_node->def(),
- kXlaControlDependenciesAttrName, &attr));
- EXPECT_EQ(attr.size(), 1);
- EXPECT_EQ(attr[0], "identity0");
-}
-
-TEST(PreprocessForEncapsulationTest, DataEdges) {
- // Build the graph:
- // "const_0" and "const_1" in host computation
- // "identityn0" = ("const_0", "const_1") in host computation 0
- // "add0" = "const_0" + "const_1" in XLA computation 0
- // "add1" = "add0" + "const_0" in XLA computation 0 & outside compilation 0
- // "identity0" = "add1" in XLA computation 0
- // "add2" = "add1" + "identity0" in host computation
- // "add3" = "add1" + "add2" in XLA computation 1
- // "add4" = "identity0" + "add2" in XLA computation 1 & outside compilation 0
- // "add5" = "identityn0"[0] + "identityn0"[1] in XLA computation 1 &
- // outside compilation 0
- // "identityn1" = ("identityn0"[0], "identityn0"[1]) in XLA computation 1 &
- // outside compilation 0
- // "identity1" = "add4" in XLA computation 1
- // "identity2" = "identity1" in host computation
- tensorflow::Scope s = tensorflow::Scope::NewRootScope();
- Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {});
- Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {});
- auto identityn0 =
- ops::IdentityN(s.WithOpName("identityn_0"), {const_0, const_1});
- Output add0 = ops::Add(s.WithOpName("add0"), const_0, const_1);
- Output add1 = ops::Add(s.WithOpName("add1"), add0, const_0);
- Output identity0 = ops::Identity(s.WithOpName("identity0"), add1);
- Output add2 = ops::Add(s.WithOpName("add2"), add1, identity0);
- Output add3 = ops::Add(s.WithOpName("add3"), add1, add2);
- Output add4 = ops::Add(s.WithOpName("add4"), identity0, add2);
- Output add5 = ops::Add(s.WithOpName("add5"), identityn0[0], identityn0[1]);
- auto identityn1 = ops::IdentityN(s.WithOpName("identityn_1"),
- {identityn0[0], identityn0[1]});
- Output identity1 = ops::Identity(s.WithOpName("identity1"), add4);
- Output identity2 = ops::Identity(s.WithOpName("identity2"), add4);
- Graph g(OpRegistry::Global());
- TF_CHECK_OK(s.ToGraph(&g));
- auto node_index = g.BuildNodeNameIndex();
-
- // Set XLA computation/outside compilation attr.
- Node *add0_node = node_index["add0"], *add1_node = node_index["add1"],
- *identity0_node = node_index["identity0"],
- *add3_node = node_index["add3"], *add4_node = node_index["add4"],
- *add5_node = node_index["add5"],
- *identityn1_node = node_index["identityn_1"],
- *identity1_node = node_index["identity1"];
- add0_node->AddAttr("_xla", "0");
- add1_node->AddAttr("_xla", "0");
- add1_node->AddAttr("_oc", "0");
- identity0_node->AddAttr("_xla", "0");
- add3_node->AddAttr("_xla", "1");
- add4_node->AddAttr("_xla", "1");
- add4_node->AddAttr("_oc", "0");
- add5_node->AddAttr("_xla", "1");
- add5_node->AddAttr("_oc", "0");
- identityn1_node->AddAttr("_xla", "1");
- identityn1_node->AddAttr("_oc", "0");
- identity1_node->AddAttr("_xla", "1");
-
- TF_CHECK_OK(PreprocessForEncapsulation(&g, "_xla", "_oc"));
-
- // Check input nodes for related data edges.
- node_index = g.BuildNodeNameIndex();
- // Step 2: add an Identity node between different XLA computations.
- Node *bridge_add1_add3 = node_index["bridge_add1_add3"];
- EXPECT_NE(bridge_add1_add3, nullptr);
- string str;
- TF_CHECK_OK(
- GetNodeAttr(bridge_add1_add3->attrs(), kBridgeSourceNodeAttrName, &str));
- EXPECT_EQ(str, "add1");
- Node *bridge_identity0_add4 = node_index["bridge_identity0_add4"];
- EXPECT_NE(bridge_identity0_add4, nullptr);
- // Step 3: add placeholder for edges between host computation and outside
- // compilation.
- EXPECT_EQ(bridge_add1_add3->def().input(0), "add1_oc_to_host_placeholder_0");
- Node *add1_oc_to_host_placeholder =
- node_index["add1_oc_to_host_placeholder_0"];
- TF_CHECK_OK(GetNodeAttr(add1_oc_to_host_placeholder->attrs(),
- kOutsideCompilationToHostOriginalNodeAttrName, &str));
- EXPECT_EQ(str, "add1");
- int i;
- TF_CHECK_OK(GetNodeAttr(add1_oc_to_host_placeholder->attrs(),
- kOutsideCompilationToHostSrcOutputAttrName, &i));
- EXPECT_EQ(i, 0);
- add4_node = node_index["add4"];
- ASSERT_NE(add4_node, nullptr);
- EXPECT_EQ(add4_node->def().input(0),
- "bridge_identity0_add4_host_to_oc_placeholder_0");
- Node *identity0_host_to_oc_placeholder =
- node_index["bridge_identity0_add4_host_to_oc_placeholder_0"];
- TF_CHECK_OK(GetNodeAttr(identity0_host_to_oc_placeholder->attrs(),
- kHostToOutsideCompilationOriginalNodeAttrName, &str));
- EXPECT_EQ(str, "bridge_identity0_add4");
- TF_CHECK_OK(GetNodeAttr(identity0_host_to_oc_placeholder->attrs(),
- kHostToOutsideCompilationSrcOutputAttrName, &i));
- EXPECT_EQ(i, 0);
-
- // Check different placeholder nodes are created for different src_output.
- Node *placeholder0 = node_index["identityn_0_host_to_oc_placeholder_0"],
- *placeholder1 = node_index["identityn_0_host_to_oc_placeholder_1"];
- EXPECT_NE(placeholder0, nullptr);
- EXPECT_NE(placeholder1, nullptr);
- // Check we only have 2 placeholder nodes created for "identityn_0".
- int placeholder_count = 0;
- for (Node *n : g.nodes()) {
- if (HasNodeAttr(n->def(), kHostToOutsideCompilationOriginalNodeAttrName)) {
- string attr;
- TF_CHECK_OK(GetNodeAttr(
- n->attrs(), kHostToOutsideCompilationOriginalNodeAttrName, &attr));
- if (attr == "identityn_0") {
- ++placeholder_count;
- }
- }
- }
- EXPECT_EQ(placeholder_count, 2);
-}
-
-TEST(PostprocessForEncapsulationTest, ControlEdges) {
- // Build the graph:
- // "const0"
- // "identity0" = "const0" (XLA computation 0)
- // "identity1" = "identity0"
- // "identity2" = "identity1" (XLA computation 1)
- // "identity3" = "identity2"
- tensorflow::Scope s = tensorflow::Scope::NewRootScope();
- Output const0 = ops::Const(s.WithOpName("const0"), 1, {});
- Output identity0 = ops::Identity(s.WithOpName("identity0"), const0);
- Output identity1 = ops::Identity(s.WithOpName("identity1"), identity0);
- Output identity2 = ops::Identity(s.WithOpName("identity2"), identity1);
- Output identity3 = ops::Identity(s.WithOpName("identity3"), identity2);
- Graph g(OpRegistry::Global());
- TF_CHECK_OK(s.ToGraph(&g));
- auto node_index = g.BuildNodeNameIndex();
-
- // Set XLA computation/outside compilation attr, and add control edges.
- Node *const0_node = node_index["const0"],
- *identity0_node = node_index["identity0"],
- *identity1_node = node_index["identity1"],
- *identity2_node = node_index["identity2"],
- *identity3_node = node_index["identity3"];
- identity1_node->AddAttr(kXlaConnectedFromOtherXlaComputationAttrName,
- std::vector<string>{"0"});
- identity1_node->AddAttr(kXlaConnectedToOtherXlaComputationAttrName,
- std::vector<string>{"1"});
- identity3_node->AddAttr(kXlaControlDependenciesAttrName,
- std::vector<string>{"const0", "identity1"});
-
- std::unordered_map<string, XlaClusterInfo> clusters;
- clusters["0"].node = identity0_node;
- clusters["1"].node = identity2_node;
- TF_CHECK_OK(PostprocessForEncapsulation(&g, "_xla", "_oc", clusters));
-
- // Case 3a: we have control edge identity0 -> identity1, and identity1 ->
- // identity2.
- bool edge_identity0_identity1 = false, edge_identity1_identity2 = false;
- for (const Edge *e : g.edges()) {
- if (!e->IsControlEdge()) {
- continue;
- }
- if (e->src() == identity0_node && e->dst() == identity1_node) {
- edge_identity0_identity1 = true;
- } else if (e->src() == identity1_node && e->dst() == identity2_node) {
- edge_identity1_identity2 = true;
- }
- }
- EXPECT_TRUE(edge_identity0_identity1);
- EXPECT_TRUE(edge_identity1_identity2);
- // Case 3b: we have control edge const0 -> identity3, and identity1 ->
- // identity3.
- bool edge_const0_identity3 = false, edge_identity1_identity3 = false;
- for (const Edge *e : g.edges()) {
- if (!e->IsControlEdge()) {
- continue;
- }
- if (e->src() == const0_node && e->dst() == identity3_node) {
- edge_const0_identity3 = true;
- } else if (e->src() == identity1_node && e->dst() == identity3_node) {
- edge_identity1_identity3 = true;
- }
- }
- EXPECT_TRUE(edge_const0_identity3);
- EXPECT_TRUE(edge_identity1_identity3);
-}
-
-TEST(PostprocessForEncapsulationTest, DataEdges) {
- // Build the graph:
- // "const0" in outside compilation "0"
- // "placeholder0" (for "const0") in host computation
- // "add0" = "placeholder0" + "placeholder0" in host computation
- // "placeholder1" (for "add0") in outside compilation 1
- // "add1" = "placeholder1" + "placeholder1" in outside compilation 1
- //
- // "bridge" = "placeholder0" in host computation
- // "placeholder2" (for "bridge") in outside compilation 1
- // "add2" = "placeholder2" + "placeholder2" in outside compilation 1
- tensorflow::Scope s = tensorflow::Scope::NewRootScope();
- Output const0 = ops::Const(s.WithOpName("const0"), 1, {});
- Output placeholder0 =
- ops::Placeholder(s.WithOpName("placeholder0"), DT_INT32);
- Output add0 = ops::Add(s.WithOpName("add0"), placeholder0, placeholder0);
- Output placeholder1 =
- ops::Placeholder(s.WithOpName("placeholder1"), DT_INT32);
- Output add1 = ops::Add(s.WithOpName("add1"), placeholder1, placeholder1);
- Output bridge = ops::Identity(s.WithOpName("bridge"), placeholder0);
- Output placeholder2 =
- ops::Placeholder(s.WithOpName("placeholder2"), DT_INT32);
- Output add2 = ops::Add(s.WithOpName("add2"), placeholder2, placeholder2);
- Graph g(OpRegistry::Global());
- TF_CHECK_OK(s.ToGraph(&g));
- auto node_index = g.BuildNodeNameIndex();
-
- // Set related attributes.
- Node *placeholder0_node = node_index["placeholder0"];
- placeholder0_node->AddAttr(kOutsideCompilationToHostOriginalNodeAttrName,
- "const0");
- placeholder0_node->AddAttr(kOutsideCompilationToHostSrcOutputAttrName, 0);
- Node *placeholder1_node = node_index["placeholder1"];
- placeholder1_node->AddAttr(kHostToOutsideCompilationOriginalNodeAttrName,
- "add0");
- placeholder1_node->AddAttr(kHostToOutsideCompilationSrcOutputAttrName, 0);
- Node *bridge_node = node_index["bridge"];
- bridge_node->AddAttr(kBridgeSourceNodeAttrName, "const0");
- Node *placeholder2_node = node_index["placeholder2"];
- placeholder2_node->AddAttr(kHostToOutsideCompilationOriginalNodeAttrName,
- "bridge");
- placeholder2_node->AddAttr(kHostToOutsideCompilationSrcOutputAttrName, 0);
-
- std::unordered_map<string, XlaClusterInfo> clusters;
- TF_CHECK_OK(PostprocessForEncapsulation(&g, "_xla", "_oc", clusters));
-
- // Result graph should be:
- // "add0" = "const0" + "const0"
- // "add1" = "add0" + "add0"
- // "add2" = "const0" + "const0"
- node_index = g.BuildNodeNameIndex();
- EXPECT_EQ(node_index.size(), 6);
- EXPECT_EQ(node_index["add0"]->def().input(0), "const0:0");
- EXPECT_EQ(node_index["add0"]->def().input(1), "const0:0");
- EXPECT_EQ(node_index["add1"]->def().input(0), "add0:0");
- EXPECT_EQ(node_index["add1"]->def().input(1), "add0:0");
- EXPECT_EQ(node_index["add2"]->def().input(0), "const0:0");
- EXPECT_EQ(node_index["add2"]->def().input(1), "const0:0");
-}
-
} // namespace tensorflow
diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass.cc b/tensorflow/compiler/jit/extract_outside_compilation_pass.cc
index e3c7e2f..baf8507 100644
--- a/tensorflow/compiler/jit/extract_outside_compilation_pass.cc
+++ b/tensorflow/compiler/jit/extract_outside_compilation_pass.cc
@@ -20,8 +20,10 @@
#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
#include "tensorflow/compiler/jit/encapsulate_util.h"
#include "tensorflow/compiler/tf2xla/dump_graph.h"
+#include "tensorflow/compiler/tf2xla/side_effect_util.h"
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
@@ -98,9 +100,12 @@
recv_at_host_builder.Attr("Toutputs", recv_at_host_dtypes);
// The correct device_ordinal will be inserted during replication in a
// subsequent rewrite.
- recv_at_host_builder.Attr("device_ordinal", 0);
+ AttrValue device_ordinal_value;
+ device_ordinal_value.set_placeholder("device_ordinal");
+ recv_at_host_builder.Attr("device_ordinal", device_ordinal_value);
recv_at_host_builder.Attr(
"key", absl::StrCat("host_compute_channel_", oc_cluster_name));
+ recv_at_host_builder.Attr(kXlaHasHostTransferAttrName, true);
recv_at_host_builder.Input(key_placeholder->name(), 0, DT_STRING);
TF_RETURN_IF_ERROR(recv_at_host_builder.Finalize(&recv_at_host_def));
Status s;
@@ -197,9 +202,12 @@
send_from_host_builder.Attr("Tinputs", send_from_host_dtypes);
// The correct device_ordinal will be inserted during replication in a
// subsequent rewrite.
- send_from_host_builder.Attr("device_ordinal", 0);
+ AttrValue device_ordinal_value;
+ device_ordinal_value.set_placeholder("device_ordinal");
+ send_from_host_builder.Attr("device_ordinal", device_ordinal_value);
send_from_host_builder.Attr(
"key", absl::StrCat("host_compute_channel_", oc_cluster_name));
+ send_from_host_builder.Attr(kXlaHasHostTransferAttrName, true);
std::vector<NodeDefBuilder::NodeOut> inputs(send_from_host_dtypes.size());
for (auto* n : ret_nodes) {
int index;
@@ -357,6 +365,47 @@
return Status::OK();
}
+// Resets "device_ordinal" attr to placeholder value for related nodes
+// (XlaRecvAtHost nodes; XlaSendFromHost nodes; If nodes containing
+// XlaRecvAtHost/XlaSendFromHost).
+Status ResetDeviceOrdinalToPlaceholderValue(Graph* g) {
+ AttrValue device_ordinal_value;
+ device_ordinal_value.set_placeholder("device_ordinal");
+ for (Node* n : g->nodes()) {
+ if (!HasNodeAttr(n->def(), kXlaHasHostTransferAttrName)) {
+ continue;
+ }
+
+ if (n->type_string() == "_XlaRecvAtHost" ||
+ n->type_string() == "_XlaSendFromHost") {
+ n->ClearAttr("device_ordinal");
+ n->AddAttr("device_ordinal", device_ordinal_value);
+ } else if (n->type_string() == "If") {
+ for (const string& attr_name :
+ std::vector<string>{"then_branch", "else_branch"}) {
+ NameAttrList branch_func;
+ TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &branch_func));
+ (*branch_func.mutable_attr())["device_ordinal"] = device_ordinal_value;
+ n->ClearAttr(attr_name);
+ n->AddAttr(attr_name, branch_func);
+ }
+ } else if (n->type_string() == "While") {
+ for (const string& attr_name : std::vector<string>{"cond", "body"}) {
+ NameAttrList branch_func;
+ TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &branch_func));
+ (*branch_func.mutable_attr())["device_ordinal"] = device_ordinal_value;
+ n->ClearAttr(attr_name);
+ n->AddAttr(attr_name, branch_func);
+ }
+ } else {
+ return errors::Internal("Unknown node marked with ",
+ kXlaHasHostTransferAttrName, ": ",
+ n->DebugString());
+ }
+ }
+ return Status::OK();
+}
+
// For an XLA computation, builds host side graph given all outside compilation
// graphs inside it. The host side graph contains:
// 1) a "sequencer" node (we will add control edge between XlaRecvAtHost and
@@ -368,8 +417,8 @@
Status ConstructHostGraph(
const string& xla_cluster_name, const string& outside_compilation_attr_name,
const std::vector<string>& outside_compilation_host_graphs,
- FunctionLibraryDefinition* fld, std::unique_ptr<Graph>* host_graph) {
- host_graph->reset(new Graph(fld));
+ FunctionLibraryDefinition* fld, const string& host_graph_func_name) {
+ Graph host_graph(fld);
// Create sequencer node in host graph.
NodeDefBuilder sequencer_builder(absl::StrCat(xla_cluster_name, "_sequencer"),
@@ -378,24 +427,34 @@
NodeDef sequencer_def;
TF_RETURN_IF_ERROR(sequencer_builder.Finalize(&sequencer_def));
Status s;
- Node* sequencer = (*host_graph)->AddNode(sequencer_def, &s);
+ Node* sequencer = host_graph.AddNode(sequencer_def, &s);
TF_RETURN_IF_ERROR(s);
// Create key placeholder in host graph.
TF_ASSIGN_OR_RETURN(
Node * key_placeholder,
- AddHostComputeKeyPlaceholder(xla_cluster_name, host_graph->get()));
+ AddHostComputeKeyPlaceholder(xla_cluster_name, &host_graph));
// For each outside compilation graph, copy them to host graph with the
// following changes:
// a) Use key_placeholder in host graph instead of its own.
- // b) Add control edge from RecvAtHost/SendFromHost to sequencer.
+ // b) Add control edge from host transfer nodes (XlaRecvAtHost,
+ // XlaSendFromHost, If/While nodes containing
+ // XlaRecvAtHost/XlaSendFromHost) to sequencer node.
// c) Clear node_def.device(), so device placer won't get confused.
for (const string& host_func : outside_compilation_host_graphs) {
VLOG(4) << "Expanding host graph " << host_func;
+ // Temporarily use "0" as "device_ordinal". It will be reset to placeholder
+ // value after we expanded all host graphs. We cannot just use placeholder
+ // value here because FunctionDef instantiation does not allow placeholder
+ // value for attributes.
+ AttrValue device_ordinal_attr;
+ device_ordinal_attr.set_i(0);
+ protobuf::Map<string, AttrValue> attrs;
+ attrs["device_ordinal"] = device_ordinal_attr;
FunctionBody* host_fbody = nullptr;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
- *fld->Find(host_func), AttrSlice(), fld,
+ *fld->Find(host_func), AttrSlice(&attrs), fld,
[&](const string& op, const OpDef** sig) {
return fld->LookUpOpDef(op, sig);
},
@@ -408,8 +467,8 @@
FixupSourceAndSinkEdges(host_fbody->graph);
std::map<const Node*, Node*> node_map;
- node_map[host_fbody->graph->source_node()] = (*host_graph)->source_node();
- node_map[host_fbody->graph->sink_node()] = (*host_graph)->sink_node();
+ node_map[host_fbody->graph->source_node()] = host_graph.source_node();
+ node_map[host_fbody->graph->sink_node()] = host_graph.sink_node();
Status s;
ReverseDFS(
*host_fbody->graph, /*enter=*/nullptr,
@@ -431,7 +490,7 @@
NodeDef copy_def = n->def();
// Change c).
copy_def.clear_device();
- copy = (*host_graph)->AddNode(copy_def, &s);
+ copy = host_graph.AddNode(copy_def, &s);
if (!s.ok()) {
return;
}
@@ -446,22 +505,23 @@
e->src()->DebugString());
return;
}
- (*host_graph)
- ->AddEdge(node_map[e->src()], e->src_output(), copy,
- e->dst_input());
+ host_graph.AddEdge(node_map[e->src()], e->src_output(), copy,
+ e->dst_input());
}
// Change b).
- if (copy->type_string() == "_XlaRecvAtHost" ||
- copy->type_string() == "_XlaSendFromHost") {
- (*host_graph)->AddControlEdge(copy, sequencer);
+ if (HasNodeAttr(copy->def(), kXlaHasHostTransferAttrName)) {
+ host_graph.AddControlEdge(copy, sequencer);
}
},
NodeComparatorID());
+
if (!s.ok()) {
return s;
}
}
+ // Reset "device_ordinal" to placeholder value.
+ TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(&host_graph));
// sequencer and key_placeholder might be dead nodes. Prune them if necessary.
// - sequencer should be pruned iff it has no input control edges from
@@ -470,21 +530,30 @@
// - key_placeholder should be pruned iff there's no RecvAtHost/SendFromHost.
// We don't need to do anything special.
if (!sequencer->in_edges().empty()) {
- (*host_graph)->AddControlEdge(sequencer, (*host_graph)->sink_node());
+ host_graph.AddControlEdge(sequencer, host_graph.sink_node());
}
PruneForReverseReachability(
- host_graph->get(),
- std::unordered_set<const Node*>{(*host_graph)->sink_node()});
+ &host_graph, std::unordered_set<const Node*>{host_graph.sink_node()});
// Postprocess edges between different outside compilations.
TF_RETURN_IF_ERROR(PostprocessEdgesBetweenOutsideCompilations(
- host_graph->get(), outside_compilation_attr_name));
+ &host_graph, outside_compilation_attr_name));
if (VLOG_IS_ON(4)) {
dump_graph::DumpGraphToFile(
absl::StrCat("extract_outside_compilation_host_graph_for_",
xla_cluster_name),
- **host_graph, fld);
+ host_graph, fld);
+ }
+
+ FunctionDef host_graph_fdef;
+ TF_RETURN_IF_ERROR(
+ GraphToFunctionDef(host_graph, host_graph_func_name, &host_graph_fdef));
+ if (fld->Find(host_graph_func_name)) {
+ TF_RETURN_IF_ERROR(
+ fld->ReplaceFunction(host_graph_func_name, host_graph_fdef));
+ } else {
+ TF_RETURN_IF_ERROR(fld->AddFunctionDef(host_graph_fdef));
}
return Status::OK();
@@ -492,8 +561,28 @@
// Expand XLA computation's outside compilation host side graph into main graph.
// Add a control edge between sequencer node and the XLA computation node.
-Status ExpandHostGraphIntoMainGraph(Graph* main_graph, Graph* host_graph,
+Status ExpandHostGraphIntoMainGraph(Graph* main_graph,
+ FunctionLibraryDefinition* fld,
+ const string& host_graph_func_name,
Node* xla_computation_node) {
+ // Temporarily use "0" as "device_ordinal". It will be rewritten with the
+ // correct value in a later pass. We cannot just use placeholder value here
+ // because FunctionDef instantiation does not allow placeholder value for
+ // attributes.
+ AttrValue device_ordinal_attr;
+ device_ordinal_attr.set_i(0);
+ protobuf::Map<string, AttrValue> attrs;
+ attrs["device_ordinal"] = device_ordinal_attr;
+ FunctionBody* fbody = nullptr;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
+ *fld->Find(host_graph_func_name), AttrSlice(&attrs), fld,
+ [&](const string& op, const OpDef** sig) {
+ return fld->LookUpOpDef(op, sig);
+ },
+ &fbody));
+ std::unique_ptr<FunctionBody> fbody_deleter(fbody);
+ Graph* host_graph = fbody->graph;
+
// We use ReverseDFS() to copy nodes. Make sure all nodes are reverse
// reachable from sink node so all nodes will be copied.
// TODO(b/77601805): consolidate copy graph functions.
@@ -545,23 +634,25 @@
return s;
}
-// Rewrites shape inference graph for outside compilation.
-// 1. If the outside compilation is a "top-level" one (not in a function of any
-// If/While/etc.), this shape inference graph might have host computation to
-// outside compilation placeholder nodes, which will cause shape inference to
-// fail. However, those nodes are not in `host_graph` any more (because we
-// have executed `PostprocessForEncapsultion`). In this case, we clear the
-// graph, and copy SendFromHost with all its predecessors from `host_graph`.
-// This case is detected by whether the SendFromHost node exists in
-// `host_graph` as well.
-// 2. Remove control edges, and prune nodes that are not useful for shape
-// inference.
+// Rewrites shape inference graph for outside compilation:
+// 1) If XlaSendFromHost also exists in `host_graph`, copy nodes from
+// `host_graph`. Because we might still have outside compilation to outside
+// compilation placeholder nodes in shape inference graph, which will prevent
+// us from inferring XlaSendFromHost shape. But in `host_graph`, we already
+// removed those placeholder nodes.
+// 2) Remove control edges.
+// 3) Prune nodes that are not useful for shape inference.
Status RewriteShapeInferenceGraph(const string& shape_inference_graph_name,
Graph* host_graph,
FunctionLibraryDefinition* fld) {
+ // Use "0" as "device_ordinal". It does not matter for shape inference.
+ AttrValue device_ordinal_attr;
+ device_ordinal_attr.set_i(0);
+ protobuf::Map<string, AttrValue> attrs;
+ attrs["device_ordinal"] = device_ordinal_attr;
FunctionBody* fbody = nullptr;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
- *fld->Find(shape_inference_graph_name), AttrSlice(), fld,
+ *fld->Find(shape_inference_graph_name), AttrSlice(&attrs), fld,
[&](const string& op, const OpDef** sig) {
return fld->LookUpOpDef(op, sig);
},
@@ -650,6 +741,7 @@
g->RemoveEdge(e);
}
}
+
// Nodes that are not reverse reachable from SendFromHost are not useful for
// shape inference. Prune them.
PruneForReverseReachability(g,
@@ -669,6 +761,567 @@
return Status::OK();
}
+// Builds XlaSendToHost node which sends cond predicate to host.
+xla::StatusOr<Node*> BuildSendIfPredNode(const string& name,
+ const string& host_transfer_key,
+ Node* pred_node, Graph* g) {
+ NodeDefBuilder send_pred_builder(name, "XlaSendToHost");
+ send_pred_builder.Attr("Tinput", DT_BOOL);
+ send_pred_builder.Attr("key", absl::StrCat(host_transfer_key, "_dtoh_0"));
+ send_pred_builder.Attr(kXlaTokenInputNodesAttrName,
+ std::vector<string>{kXlaTokenArgNodeName});
+ send_pred_builder.Input(pred_node->name(), 0, DT_BOOL);
+ NodeDef send_pred_def;
+ TF_RETURN_IF_ERROR(send_pred_builder.Finalize(&send_pred_def));
+ Status s;
+ Node* send_pred_node = g->AddNode(send_pred_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ g->AddEdge(pred_node, 0, send_pred_node, 0);
+ return send_pred_node;
+}
+
+// Replaces key placeholder node with an _Arg node.
+Status ReplaceKeyPlaceholderWithArgNode(const string& xla_cluster_name,
+ const string& func_name,
+ FunctionLibraryDefinition* fld) {
+ // Temporarily use "0" as "device_ordinal". It will be reset to placeholder
+ // value after rewriting.
+ AttrValue device_ordinal_attr;
+ device_ordinal_attr.set_i(0);
+ protobuf::Map<string, AttrValue> attrs;
+ attrs["device_ordinal"] = device_ordinal_attr;
+ FunctionBody* fbody = nullptr;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
+ *fld->Find(func_name), AttrSlice(&attrs), fld,
+ [&](const string& op, const OpDef** sig) {
+ return fld->LookUpOpDef(op, sig);
+ },
+ &fbody));
+ std::unique_ptr<FunctionBody> fbody_deleter(fbody);
+ Graph* g = fbody->graph;
+
+ // Find or create the key placeholder node.
+ Node* key_placeholder = nullptr;
+ for (Node* n : g->nodes()) {
+ if (IsKeyPlaceholderNode(*n)) {
+ key_placeholder = n;
+ break;
+ }
+ }
+ if (!key_placeholder) {
+ TF_ASSIGN_OR_RETURN(key_placeholder,
+ AddHostComputeKeyPlaceholder(xla_cluster_name, g));
+ }
+
+ // Build the _Arg node, and replace key placeholder node with it.
+ NodeDefBuilder arg_builder("key_arg", FunctionLibraryDefinition::kArgOp);
+ arg_builder.Attr("T", DT_STRING);
+ arg_builder.Attr("index", 0);
+ NodeDef arg_def;
+ TF_RETURN_IF_ERROR(arg_builder.Finalize(&arg_def));
+ TF_RETURN_IF_ERROR(ReplaceNode(g, key_placeholder, arg_def).status());
+
+ // Reset "device_ordinal" to placeholder value.
+ TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(g));
+
+ FunctionDef replace_fdef;
+ TF_RETURN_IF_ERROR(GraphToFunctionDef(*g, func_name, &replace_fdef));
+ TF_RETURN_IF_ERROR(fld->ReplaceFunction(func_name, replace_fdef));
+ return Status::OK();
+}
+
+// Builds host side graph for If node.
+Status BuildHostGraphForIfNode(const string& xla_cluster_attr_name,
+ const string& outside_compilation_attr_name,
+ const string& xla_cluster_name,
+ const string& if_node_name,
+ const string& host_transfer_key,
+ const string& host_graph_func_name,
+ FunctionLibraryDefinition* fld,
+ const string& then_branch_host_func_name,
+ const string& else_branch_host_func_name) {
+ Graph host_graph(fld);
+ string outside_compilation_name = absl::StrCat("oc_if_", if_node_name);
+ AttrValue device_ordinal_value;
+ device_ordinal_value.set_placeholder("device_ordinal");
+
+ // Step 1: add key placeholder node.
+ TF_ASSIGN_OR_RETURN(
+ Node * key_placeholder,
+ AddHostComputeKeyPlaceholder(xla_cluster_name, &host_graph));
+
+ // Step 2: build XlaRecvAtHost node to recv predicate.
+ NodeDefBuilder recv_pred_builder(
+ absl::StrCat("recv_oc_if_pred_", if_node_name), "_XlaRecvAtHost");
+ recv_pred_builder.Attr("Toutputs", std::vector<DataType>{DT_BOOL});
+ recv_pred_builder.Attr("key", host_transfer_key);
+ recv_pred_builder.Attr("device_ordinal", device_ordinal_value);
+ recv_pred_builder.Attr(xla_cluster_attr_name, xla_cluster_name);
+ recv_pred_builder.Attr(outside_compilation_attr_name,
+ outside_compilation_name);
+ recv_pred_builder.Attr(kXlaHasHostTransferAttrName, true);
+ recv_pred_builder.Input(key_placeholder->name(), 0, DT_STRING);
+ NodeDef recv_pred_def;
+ TF_RETURN_IF_ERROR(recv_pred_builder.Finalize(&recv_pred_def));
+ Status s;
+ Node* recv_pred_node = host_graph.AddNode(recv_pred_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ host_graph.AddEdge(key_placeholder, 0, recv_pred_node, 0);
+
+ // Step 3: rewrite `{then, else}_branch_host_func_name`, replace key
+ // placeholder with an _Arg node.
+ TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode(
+ xla_cluster_name, then_branch_host_func_name, fld));
+ TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode(
+ xla_cluster_name, else_branch_host_func_name, fld));
+
+ // Step 4: build If node to choose between `{then, else}_branch_host_graph`.
+ NodeDefBuilder if_builder(absl::StrCat("oc_if_", if_node_name), "If");
+ if_builder.Attr("Tcond", DT_BOOL);
+ if_builder.Attr("Tin", std::vector<DataType>{DT_STRING});
+ if_builder.Attr("Tout", std::vector<DataType>{});
+ NameAttrList host_then_branch, host_else_branch;
+ host_then_branch.set_name(then_branch_host_func_name);
+ (*host_then_branch.mutable_attr())["device_ordinal"] = device_ordinal_value;
+ host_else_branch.set_name(else_branch_host_func_name);
+ (*host_else_branch.mutable_attr())["device_ordinal"] = device_ordinal_value;
+ if_builder.Attr("then_branch", host_then_branch);
+ if_builder.Attr("else_branch", host_else_branch);
+ if_builder.Attr(kXlaHasHostTransferAttrName, true);
+ if_builder.Attr(xla_cluster_attr_name, xla_cluster_name);
+ if_builder.Attr(outside_compilation_attr_name, outside_compilation_name);
+ if_builder.Input(recv_pred_node->name(), 0, DT_BOOL);
+ std::vector<NodeDefBuilder::NodeOut> if_inputs{
+ {key_placeholder->name(), 0, DT_STRING}};
+ if_builder.Input(if_inputs);
+ NodeDef if_def;
+ TF_RETURN_IF_ERROR(if_builder.Finalize(&if_def));
+ Node* if_node = host_graph.AddNode(if_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ host_graph.AddEdge(recv_pred_node, 0, if_node, 0);
+ host_graph.AddEdge(key_placeholder, 0, if_node, 1);
+
+ // Convert `host_graph` to function, and add a "device_ordinal" attr.
+ FunctionDef oc_host_graph_fdef;
+ TF_RETURN_IF_ERROR(GraphToFunctionDef(host_graph, host_graph_func_name,
+ &oc_host_graph_fdef));
+ if (fld->Find(host_graph_func_name)) {
+ TF_RETURN_IF_ERROR(
+ fld->ReplaceFunction(host_graph_func_name, oc_host_graph_fdef));
+ } else {
+ TF_RETURN_IF_ERROR(fld->AddFunctionDef(oc_host_graph_fdef));
+ }
+
+ return Status::OK();
+}
+
+// Rewrites loop cond to add a node which sends loop cond to host.
+Status AddSendLoopPredToLoopCond(FunctionLibraryDefinition* fld,
+ const NameAttrList& loop_cond_func,
+ const string& while_node_name,
+ const string& host_transfer_key) {
+ // Instantiate the loop cond function.
+ FunctionBody* fbody = nullptr;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
+ *fld->Find(loop_cond_func.name()), AttrSlice(&loop_cond_func.attr()), fld,
+ [&](const string& op, const OpDef** sig) {
+ return fld->LookUpOpDef(op, sig);
+ },
+ &fbody));
+ std::unique_ptr<FunctionBody> fbody_deleter(fbody);
+ Graph* g = fbody->graph;
+
+ // Find the _Retval node and the loop cond node.
+ Node* ret_node = nullptr;
+ for (Node* n : g->nodes()) {
+ if (n->type_string() == "_Retval") {
+ if (ret_node) {
+ return errors::Internal("Multiple return node for loop cond function ",
+ loop_cond_func.name(), ": ",
+ ret_node->DebugString(), " and ",
+ n->DebugString());
+ } else {
+ ret_node = n;
+ }
+ }
+ }
+ if (!ret_node) {
+ return errors::Internal("No _Retval node for loop cond function ",
+ loop_cond_func.name());
+ }
+ Node* loop_cond;
+ TF_RETURN_IF_ERROR(ret_node->input_node(0, &loop_cond));
+
+ // Build the XlaSendToHost node.
+ NodeDefBuilder send_loop_cond_builder(
+ absl::StrCat("send_oc_while_cond_", while_node_name), "XlaSendToHost");
+ send_loop_cond_builder.Attr("Tinput", DT_BOOL);
+ send_loop_cond_builder.Attr("key",
+ absl::StrCat(host_transfer_key, "_dtoh_0"));
+ send_loop_cond_builder.Attr(kXlaTokenInputNodesAttrName,
+ std::vector<string>{kXlaTokenArgNodeName});
+ send_loop_cond_builder.Input(loop_cond->name(), 0, DT_BOOL);
+ NodeDef send_loop_cond_def;
+ TF_RETURN_IF_ERROR(send_loop_cond_builder.Finalize(&send_loop_cond_def));
+ Status s;
+ Node* send_loop_cond_node = g->AddNode(send_loop_cond_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ g->AddEdge(loop_cond, 0, send_loop_cond_node, 0);
+
+ // Replace original function.
+ FunctionDef replace_fdef;
+ TF_RETURN_IF_ERROR(
+ GraphToFunctionDef(*g, loop_cond_func.name(), &replace_fdef));
+ TF_RETURN_IF_ERROR(fld->ReplaceFunction(loop_cond_func.name(), replace_fdef));
+
+ return Status::OK();
+}
+
+// Rewrites while loop cond function for host.
+Status RewriteHostWhileLoopCond(
+ const string& cond_host_func_name, const string& while_node_name,
+ const string& host_transfer_key, const string& xla_cluster_attr_name,
+ const string& xla_cluster_name, const string& outside_compilation_attr_name,
+ const string& outside_compilation_name, FunctionLibraryDefinition* fld) {
+ // Replace key placeholder node with _Arg node.
+ TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode(
+ xla_cluster_name, cond_host_func_name, fld));
+
+ // Instantiate cond function.
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> attrs;
+ attrs["device_ordinal"] = device_ordinal_temp_value;
+ FunctionBody* cond_fbody = nullptr;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
+ *fld->Find(cond_host_func_name), AttrSlice(&attrs), fld,
+ [&](const string& op, const OpDef** sig) {
+ return fld->LookUpOpDef(op, sig);
+ },
+ &cond_fbody));
+ std::unique_ptr<FunctionBody> cond_fbody_deleter(cond_fbody);
+ Graph* cond_graph = cond_fbody->graph;
+ Node* key_arg = nullptr;
+ for (Node* n : cond_graph->nodes()) {
+ if (n->type_string() == "_Arg") {
+ key_arg = n;
+ }
+ }
+ if (!key_arg) {
+ return errors::Internal(
+ "No _Arg node found for host compute key in function ",
+ cond_host_func_name);
+ }
+
+ // Add an XlaRecvAtHost node to use as cond function return value.
+ // We don't need to set kXlaHasHostTransferAttrName for this node, because
+ // it's already added for the "While" node on the host.
+ NodeDefBuilder recv_pred_builder(
+ absl::StrCat("recv_oc_while_cond_", while_node_name), "_XlaRecvAtHost");
+ recv_pred_builder.Attr("Toutputs", std::vector<DataType>{DT_BOOL});
+ recv_pred_builder.Attr("key", host_transfer_key);
+ AttrValue device_ordinal_value;
+ device_ordinal_value.set_placeholder("device_ordinal");
+ recv_pred_builder.Attr("device_ordinal", device_ordinal_value);
+ recv_pred_builder.Attr(xla_cluster_attr_name, xla_cluster_name);
+ recv_pred_builder.Attr(outside_compilation_attr_name,
+ outside_compilation_name);
+ recv_pred_builder.Input(key_arg->name(), 0, DT_STRING);
+ NodeDef recv_pred_def;
+ TF_RETURN_IF_ERROR(recv_pred_builder.Finalize(&recv_pred_def));
+ Status s;
+ Node* recv_pred_node = cond_graph->AddNode(recv_pred_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ cond_graph->AddEdge(key_arg, 0, recv_pred_node, 0);
+ NodeDefBuilder ret_builder(
+ absl::StrCat("recv_oc_while_cond_ret_", while_node_name), "_Retval");
+ ret_builder.Attr("T", DT_BOOL);
+ ret_builder.Attr("index", 0);
+ ret_builder.Input(recv_pred_node->name(), 0, DT_BOOL);
+ NodeDef ret_def;
+ TF_RETURN_IF_ERROR(ret_builder.Finalize(&ret_def));
+ Node* ret_node = cond_graph->AddNode(ret_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ cond_graph->AddEdge(recv_pred_node, 0, ret_node, 0);
+
+ // Reset device_ordinal to placeholder value.
+ TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(cond_graph));
+
+ // Replace original function.
+ FunctionDef cond_replace_fdef;
+ TF_RETURN_IF_ERROR(
+ GraphToFunctionDef(*cond_graph, cond_host_func_name, &cond_replace_fdef));
+ TF_RETURN_IF_ERROR(
+ fld->ReplaceFunction(cond_host_func_name, cond_replace_fdef));
+
+ return Status::OK();
+}
+
+// Rewrites while loop body function for host.
+Status RewriteHostWhileLoopBody(
+ const string& body_host_func_name, const string& while_node_name,
+ const string& host_transfer_key, const string& xla_cluster_attr_name,
+ const string& xla_cluster_name, const string& outside_compilation_attr_name,
+ const string& outside_compilation_name, FunctionLibraryDefinition* fld) {
+ // Replace key placeholder node with _Arg node.
+ TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode(
+ xla_cluster_name, body_host_func_name, fld));
+
+ // Instantiate body function.
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> attrs;
+ attrs["device_ordinal"] = device_ordinal_temp_value;
+ FunctionBody* body_fbody = nullptr;
+ TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
+ *fld->Find(body_host_func_name), AttrSlice(&attrs), fld,
+ [&](const string& op, const OpDef** sig) {
+ return fld->LookUpOpDef(op, sig);
+ },
+ &body_fbody));
+ std::unique_ptr<FunctionBody> body_fbody_deleter(body_fbody);
+ Graph* body_graph = body_fbody->graph;
+ Node* key_arg = nullptr;
+ for (Node* n : body_graph->nodes()) {
+ if (n->type_string() == "_Arg") {
+ key_arg = n;
+ }
+ }
+ if (!key_arg) {
+ return errors::Internal(
+ "No _Arg node found for host compute key in function ",
+ body_host_func_name);
+ }
+
+ // Add a _Retval node to loop body.
+ NodeDefBuilder ret_builder(
+ absl::StrCat("recv_oc_while_body_ret_", while_node_name), "_Retval");
+ ret_builder.Attr("T", DT_STRING);
+ ret_builder.Attr("index", 0);
+ ret_builder.Input(key_arg->name(), 0, DT_STRING);
+ NodeDef ret_def;
+ TF_RETURN_IF_ERROR(ret_builder.Finalize(&ret_def));
+ Status s;
+ Node* ret_node = body_graph->AddNode(ret_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ body_graph->AddEdge(key_arg, 0, ret_node, 0);
+
+ // Reset device_ordinal to placeholder value.
+ TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(body_graph));
+
+ // Replace original function.
+ FunctionDef body_replace_fdef;
+ TF_RETURN_IF_ERROR(
+ GraphToFunctionDef(*body_graph, body_host_func_name, &body_replace_fdef));
+ TF_RETURN_IF_ERROR(
+ fld->ReplaceFunction(body_host_func_name, body_replace_fdef));
+
+ return Status::OK();
+}
+
+// Builds host side graph for while node.
+Status BuildHostGraphForWhileNode(
+ const string& xla_cluster_attr_name,
+ const string& outside_compilation_attr_name, const string& xla_cluster_name,
+ const string& while_node_name, const string& host_transfer_key,
+ const string& host_graph_func_name, FunctionLibraryDefinition* fld,
+ const string& cond_host_func_name, const string& body_host_func_name) {
+ Graph host_graph(fld);
+ string outside_compilation_name = absl::StrCat("oc_while_", while_node_name);
+
+ // Step 1: add key placeholder node.
+ TF_ASSIGN_OR_RETURN(
+ Node * key_placeholder,
+ AddHostComputeKeyPlaceholder(xla_cluster_name, &host_graph));
+
+ // Step 2: rewrite cond function.
+ TF_RETURN_IF_ERROR(RewriteHostWhileLoopCond(
+ cond_host_func_name, while_node_name, host_transfer_key,
+ xla_cluster_attr_name, xla_cluster_name, outside_compilation_attr_name,
+ outside_compilation_name, fld));
+
+ // Step 3: rewrite body function.
+ TF_RETURN_IF_ERROR(RewriteHostWhileLoopBody(
+ body_host_func_name, while_node_name, host_transfer_key,
+ xla_cluster_attr_name, xla_cluster_name, outside_compilation_attr_name,
+ outside_compilation_name, fld));
+
+ // Step 4: build While node.
+ NodeDefBuilder while_builder(absl::StrCat("oc_while_", while_node_name),
+ "While");
+ while_builder.Attr("T", std::vector<DataType>{DT_STRING});
+ NameAttrList func;
+ AttrValue device_ordinal_value;
+ device_ordinal_value.set_placeholder("device_ordinal");
+ (*func.mutable_attr())["device_ordinal"] = device_ordinal_value;
+ func.set_name(cond_host_func_name);
+ while_builder.Attr("cond", func);
+ func.set_name(body_host_func_name);
+ while_builder.Attr("body", func);
+ while_builder.Attr(kXlaHasHostTransferAttrName, true);
+ while_builder.Attr(xla_cluster_attr_name, xla_cluster_name);
+ while_builder.Attr(outside_compilation_attr_name, outside_compilation_name);
+ std::vector<NodeDefBuilder::NodeOut> while_inputs{
+ {key_placeholder->name(), 0, DT_STRING}};
+ while_builder.Input(while_inputs);
+ NodeDef while_def;
+ TF_RETURN_IF_ERROR(while_builder.Finalize(&while_def));
+ Status s;
+ Node* while_node = host_graph.AddNode(while_def, &s);
+ TF_RETURN_IF_ERROR(s);
+ host_graph.AddEdge(key_placeholder, 0, while_node, 0);
+
+ // Convert `host_graph` to function.
+ FunctionDef oc_host_graph_fdef;
+ TF_RETURN_IF_ERROR(GraphToFunctionDef(host_graph, host_graph_func_name,
+ &oc_host_graph_fdef));
+ if (fld->Find(host_graph_func_name)) {
+ TF_RETURN_IF_ERROR(
+ fld->ReplaceFunction(host_graph_func_name, oc_host_graph_fdef));
+ } else {
+ TF_RETURN_IF_ERROR(fld->AddFunctionDef(oc_host_graph_fdef));
+ }
+
+ return Status::OK();
+}
+
+Status ExtractOutsideCompilationForNodesWithAssociatedFunctions(
+ Graph* g, const string& xla_cluster_attr_name,
+ const string& outside_compilation_attr_name, const string& xla_cluster_name,
+ const std::map<string, int>& host_compute_core,
+ FunctionLibraryDefinition* fld, std::vector<string>* host_graphs,
+ std::vector<string>* shape_inference_graphs,
+ bool* has_outside_compilation) {
+ std::vector<Node*> if_nodes, while_nodes;
+ for (Node* n : g->nodes()) {
+ if (n->type_string() == "If") {
+ if_nodes.push_back(n);
+ } else if (n->type_string() == "While") {
+ while_nodes.push_back(n);
+ }
+ }
+
+ for (Node* n : if_nodes) {
+ // Instantiate "then_branch" and "else_branch".
+ NameAttrList then_branch, else_branch;
+ TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "then_branch", &then_branch));
+ TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "else_branch", &else_branch));
+
+ // Extract outside compilation for then_branch and else_branch.
+ bool then_branch_has_outside_compilation = false;
+ bool else_branch_has_outside_compilation = false;
+ string then_branch_host_func_name =
+ absl::StrCat("oc_then_branch_host_if_", n->name()),
+ else_branch_host_func_name =
+ absl::StrCat("oc_else_branch_host_if_", n->name());
+ string then_branch_xla_func_name = absl::StrCat(then_branch.name(), "_oc"),
+ else_branch_xla_func_name = absl::StrCat(else_branch.name(), "_oc");
+ TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction(
+ xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
+ then_branch, then_branch_xla_func_name, then_branch_host_func_name,
+ host_compute_core, fld, shape_inference_graphs,
+ &then_branch_has_outside_compilation));
+ TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction(
+ xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
+ else_branch, else_branch_xla_func_name, else_branch_host_func_name,
+ host_compute_core, fld, shape_inference_graphs,
+ &else_branch_has_outside_compilation));
+
+ // If then/else branch do not have outside compilation, nothing to do.
+ if (!then_branch_has_outside_compilation &&
+ !else_branch_has_outside_compilation) {
+ continue;
+ }
+
+ *has_outside_compilation = true;
+
+ // Change If node to call the new functions.
+ then_branch.set_name(then_branch_xla_func_name);
+ n->ClearAttr("then_branch");
+ n->AddAttr("then_branch", then_branch);
+ else_branch.set_name(else_branch_xla_func_name);
+ n->ClearAttr("else_branch");
+ n->AddAttr("else_branch", else_branch);
+
+ string host_transfer_key = absl::StrCat("oc_if_pred_", n->name());
+
+ // XLA computation: add a SendToHost node to send cond predicate.
+ Node* pred_node;
+ TF_RETURN_IF_ERROR(n->input_node(0, &pred_node));
+ TF_ASSIGN_OR_RETURN(
+ Node * send_pred_node,
+ BuildSendIfPredNode(absl::StrCat("send_oc_if_pred_", n->name()),
+ host_transfer_key, pred_node, g));
+ n->AddAttr(kXlaTokenInputNodesAttrName,
+ std::vector<string>{send_pred_node->name()});
+
+ // Build host side graph for the "If" node.
+ string oc_host_graph_name = absl::StrCat("oc_if_host_graph_", n->name());
+ TF_RETURN_IF_ERROR(BuildHostGraphForIfNode(
+ xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
+ n->name(), host_transfer_key, oc_host_graph_name, fld,
+ then_branch_host_func_name, else_branch_host_func_name));
+ host_graphs->push_back(oc_host_graph_name);
+ }
+
+ for (Node* n : while_nodes) {
+ // Instantiate "cond" and "body".
+ NameAttrList cond, body;
+ TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "cond", &cond));
+ TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "body", &body));
+
+ // Extract outside compilation for cond and body.
+ bool cond_has_outside_compilation = false;
+ bool body_has_outside_compilation = false;
+ string cond_host_func_name = absl::StrCat("oc_cond_host_while_", n->name()),
+ body_host_func_name = absl::StrCat("oc_body_host_while_", n->name());
+ string cond_xla_func_name = absl::StrCat(cond.name(), "_oc"),
+ body_xla_func_name = absl::StrCat(body.name(), "_oc");
+ TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction(
+ xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
+ cond, cond_xla_func_name, cond_host_func_name, host_compute_core, fld,
+ shape_inference_graphs, &cond_has_outside_compilation));
+ TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction(
+ xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
+ body, body_xla_func_name, body_host_func_name, host_compute_core, fld,
+ shape_inference_graphs, &body_has_outside_compilation));
+
+ // If cond/body do not have outside compilation, nothing to do.
+ if (!cond_has_outside_compilation && !body_has_outside_compilation) {
+ continue;
+ }
+
+ *has_outside_compilation = true;
+
+ // Change While node to call the new functions.
+ cond.set_name(cond_xla_func_name);
+ n->ClearAttr("cond");
+ n->AddAttr("cond", cond);
+ body.set_name(body_xla_func_name);
+ n->ClearAttr("body");
+ n->AddAttr("body", body);
+
+ string host_transfer_key = absl::StrCat("oc_while_pred_", n->name());
+
+ // XLA computation: rewrite cond function to add a SendToHost node to send
+ // loop predicate.
+ TF_RETURN_IF_ERROR(
+ AddSendLoopPredToLoopCond(fld, cond, n->name(), host_transfer_key));
+ n->AddAttr(kXlaTokenInputNodesAttrName,
+ std::vector<string>{kXlaTokenArgNodeName});
+
+ // Build host side graph for the "While" node.
+ string oc_host_graph_name = absl::StrCat("oc_while_host_graph_", n->name());
+ TF_RETURN_IF_ERROR(BuildHostGraphForWhileNode(
+ xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
+ n->name(), host_transfer_key, oc_host_graph_name, fld,
+ cond_host_func_name, body_host_func_name));
+ host_graphs->push_back(oc_host_graph_name);
+ }
+
+ return Status::OK();
+}
+
} // namespace
Status RewriteOutsideCompilationSubgraphFn::operator()(
@@ -755,12 +1408,15 @@
// it with HostCompute node later.
AddNodeAttr("_outside_compilation_subgraph", old_name, node_def);
if (shapes) {
- AddNodeAttr("shape_inference_graph", "", node_def);
+ NameAttrList shape_inference_graph;
+ AddNodeAttr("shape_inference_graph", shape_inference_graph, node_def);
AddNodeAttr("shapes", *shapes, node_def);
} else {
string shape_inference_func_name =
absl::StrCat("_outside_compilation_shape_inference_", new_name);
- AddNodeAttr("shape_inference_graph", shape_inference_func_name, node_def);
+ NameAttrList shape_inference_graph;
+ shape_inference_graph.set_name(shape_inference_func_name);
+ AddNodeAttr("shape_inference_graph", shape_inference_graph, node_def);
AddNodeAttr("shapes", std::vector<TensorShapeProto>{}, node_def);
}
AddNodeAttr("ancestors", std::vector<string>{}, node_def);
@@ -775,11 +1431,10 @@
const string& xla_cluster_attr_name,
const string& outside_compilation_attr_name, const string& xla_cluster_name,
const NameAttrList& func_name_attrs, const string& new_func_name,
+ const string& host_graph_func_name,
const std::map<string, int>& host_compute_core,
- FunctionLibraryDefinition* fld, std::unique_ptr<Graph>* host_graph,
- std::vector<string>* shape_inference_graphs,
+ FunctionLibraryDefinition* fld, std::vector<string>* shape_inference_graphs,
bool* has_outside_compilation) {
- // Early return if function does not have any outside compilation nodes.
const string& func_name = func_name_attrs.name();
const FunctionDef* fdef = fld->Find(func_name);
if (!fdef) {
@@ -792,9 +1447,8 @@
break;
}
}
- if (!has_outside_compilation) {
- return Status::OK();
- }
+ // We cannot early return here, because we might have outside compilation in
+ // If/While function body.
// Convert the function to graph.
FunctionBody* fbody = nullptr;
@@ -835,11 +1489,11 @@
// If we could not infer shapes for XlaSendFromHost inputs statically, we
// will set the "shape_inference_graph" attribute. In that case, copy
// outside compilation subgraph as shape inference graph in `fld`.
- string shape_inference_graph;
+ NameAttrList shape_inference_graph;
TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "shape_inference_graph",
&shape_inference_graph));
- if (!shape_inference_graph.empty()) {
- shape_inference_graphs->push_back(shape_inference_graph);
+ if (!shape_inference_graph.name().empty()) {
+ shape_inference_graphs->push_back(shape_inference_graph.name());
const FunctionDef* xla_fdef = fld->Find(n->name());
if (!xla_fdef) {
@@ -847,9 +1501,9 @@
}
FunctionDef shape_inference_fdef = *xla_fdef;
shape_inference_fdef.mutable_signature()->set_name(
- shape_inference_graph);
- if (fld->Find(shape_inference_graph)) {
- TF_RETURN_IF_ERROR(fld->ReplaceFunction(shape_inference_graph,
+ shape_inference_graph.name());
+ if (fld->Find(shape_inference_graph.name())) {
+ TF_RETURN_IF_ERROR(fld->ReplaceFunction(shape_inference_graph.name(),
shape_inference_fdef));
} else {
TF_RETURN_IF_ERROR(fld->AddFunctionDef(shape_inference_fdef));
@@ -867,12 +1521,17 @@
*graph_out, fld);
}
+ // Handle nodes with associated functions.
+ TF_RETURN_IF_ERROR(ExtractOutsideCompilationForNodesWithAssociatedFunctions(
+ graph_out.get(), xla_cluster_attr_name, outside_compilation_attr_name,
+ xla_cluster_name, host_compute_core, fld,
+ &outside_compilation_host_graphs, shape_inference_graphs,
+ has_outside_compilation));
+
// Construct host graph.
- if (!outside_compilation_host_graphs.empty()) {
- TF_RETURN_IF_ERROR(
- ConstructHostGraph(xla_cluster_name, outside_compilation_attr_name,
- outside_compilation_host_graphs, fld, host_graph));
- }
+ TF_RETURN_IF_ERROR(ConstructHostGraph(
+ xla_cluster_name, outside_compilation_attr_name,
+ outside_compilation_host_graphs, fld, host_graph_func_name));
// Remove the outside compilation graphs from function library.
for (const string& func : outside_compilation_host_graphs) {
@@ -909,24 +1568,17 @@
auto const& host_compute_core = iter.second.host_compute_core;
bool has_outside_compilation;
- std::unique_ptr<Graph> host_graph;
+ string host_graph_func_name = absl::StrCat("oc_host_graph_", n->name());
TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction(
xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name,
- func_name_attrs, func_name_attrs.name(), host_compute_core, fld,
- &host_graph, &shape_inference_graphs, &has_outside_compilation));
- if (host_graph) {
- TF_RETURN_IF_ERROR(ExpandHostGraphIntoMainGraph(g, host_graph.get(), n));
- }
+ func_name_attrs, func_name_attrs.name(), host_graph_func_name,
+ host_compute_core, fld, &shape_inference_graphs,
+ &has_outside_compilation));
+ TF_RETURN_IF_ERROR(
+ ExpandHostGraphIntoMainGraph(g, fld, host_graph_func_name, n));
+ TF_RETURN_IF_ERROR(fld->RemoveFunction(host_graph_func_name));
}
- if (VLOG_IS_ON(4)) {
- dump_graph::DumpGraphToFile("extract_outside_compilation_expanded", *g,
- fld);
- }
-
- TF_RETURN_IF_ERROR(PostprocessForEncapsulation(
- g, xla_cluster_attr_name, outside_compilation_attr_name, clusters));
-
for (auto shape_inference_graph_name : shape_inference_graphs) {
TF_RETURN_IF_ERROR(
RewriteShapeInferenceGraph(shape_inference_graph_name, g, fld));
diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass.h b/tensorflow/compiler/jit/extract_outside_compilation_pass.h
index 2a4f07c..e07e7c5 100644
--- a/tensorflow/compiler/jit/extract_outside_compilation_pass.h
+++ b/tensorflow/compiler/jit/extract_outside_compilation_pass.h
@@ -88,9 +88,10 @@
const string& xla_cluster_attr_name,
const string& outside_compilation_attr_name, const string& xla_cluster_name,
const NameAttrList& func_name_attrs, const string& new_func_name,
+ const string& host_graph_func_name,
const std::map<string, int>& host_compute_core,
- FunctionLibraryDefinition* fld, std::unique_ptr<Graph>* host_graph,
- std::vector<string>* shape_inference_graphs, bool* has_outside_compilation);
+ FunctionLibraryDefinition* fld, std::vector<string>* shape_inference_graphs,
+ bool* has_outside_compilation);
// Rewrites XLA computation in `clusters` to replace outside compilation nodes
// with XlaHostCompute, and moves those outside compilations into `g`. If shapes
diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc b/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc
index bff9561..0887fbc 100644
--- a/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc
+++ b/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc
@@ -19,8 +19,10 @@
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/function_ops.h"
+#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/encapsulate_util.h"
+#include "tensorflow/compiler/xla/test.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/function.h"
@@ -109,10 +111,10 @@
}
EXPECT_TRUE(has_control_edge_to_send_from_host);
// Verify step 7: necessary attrs added to call_node_def.
- string shape_inference_graph;
+ NameAttrList shape_inference_graph;
TF_CHECK_OK(GetNodeAttr(AttrSlice(&call_node_def.attr()),
"shape_inference_graph", &shape_inference_graph));
- EXPECT_EQ(shape_inference_graph,
+ EXPECT_EQ(shape_inference_graph.name(),
"_outside_compilation_shape_inference_cluster_0");
}
@@ -249,27 +251,26 @@
protobuf::Map<string, tensorflow::AttrValue> attrs;
std::map<string, int> host_compute_core = {{"0", 1}, {"1", 0}};
- std::unique_ptr<Graph> host_graph;
std::vector<string> shape_inference_graphs;
bool has_outside_compilation;
NameAttrList name_attrs;
name_attrs.set_name("cluster");
*name_attrs.mutable_attr() = attrs;
TF_CHECK_OK(ExtractOutsideCompilationForFunction(
- "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten",
- host_compute_core, &fld, &host_graph, &shape_inference_graphs,
+ "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
+ host_compute_core, &fld, &shape_inference_graphs,
&has_outside_compilation));
// Get rewritten XLA computation function.
- FunctionBody *fbody = nullptr;
- TF_CHECK_OK(FunctionDefToBodyHelper(*fld.Find("cluster_rewritten"),
- AttrSlice(), &fld,
- [&](const string &op, const OpDef **sig) {
- return fld.LookUpOpDef(op, sig);
- },
- &fbody));
- std::unique_ptr<FunctionBody> fbody_deleter(fbody);
- auto node_name_index = fbody->graph->BuildNodeNameIndex();
+ FunctionBody *xla_fbody = nullptr;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("cluster_rewritten"), AttrSlice(), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &xla_fbody));
+ std::unique_ptr<FunctionBody> xla_fbody_deleter(xla_fbody);
+ auto node_name_index = xla_fbody->graph->BuildNodeNameIndex();
// Check XlaHostCompute nodes.
Node *host_compute_0 = node_name_index["outside_compilation_0_host_compute"];
@@ -292,18 +293,31 @@
EXPECT_EQ(shapes[0].dim_size(), 1);
// Check XlaHostCompute nodes' "shape_inference_graph" attr. Both should have
// empty values.
- string shape_inference_graph;
+ NameAttrList shape_inference_graph;
TF_CHECK_OK(GetNodeAttr(host_compute_0->attrs(), "shape_inference_graph",
&shape_inference_graph));
- EXPECT_EQ(shape_inference_graph, "");
+ EXPECT_EQ(shape_inference_graph.name(), "");
TF_CHECK_OK(GetNodeAttr(host_compute_1->attrs(), "shape_inference_graph",
&shape_inference_graph));
- EXPECT_EQ(shape_inference_graph, "");
+ EXPECT_EQ(shape_inference_graph.name(), "");
// Check `shape_inference_graphs`.
EXPECT_EQ(shape_inference_graphs.size(), 0);
- // Check `host_graph`: verify we have key placeholder and sequencer.
+ // Check host graph: verify we have key placeholder and sequencer.
+ FunctionBody *host_fbody = nullptr;
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> host_func_attrs;
+ host_func_attrs["device_ordinal"] = device_ordinal_temp_value;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &host_fbody));
+ std::unique_ptr<FunctionBody> host_fbody_deleter(host_fbody);
+ Graph *host_graph = host_fbody->graph;
Node *key_placeholder = nullptr, *sequencer = nullptr;
for (Node *n : host_graph->nodes()) {
if (n->type_string() == "Placeholder" &&
@@ -365,25 +379,37 @@
protobuf::Map<string, tensorflow::AttrValue> attrs;
std::map<string, int> host_compute_core = {{"0", 1}, {"1", 0}};
- std::unique_ptr<Graph> host_graph;
std::vector<string> shape_inference_graphs;
bool has_outside_compilation;
NameAttrList name_attrs;
name_attrs.set_name("cluster");
*name_attrs.mutable_attr() = attrs;
TF_CHECK_OK(ExtractOutsideCompilationForFunction(
- "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten",
- host_compute_core, &fld, &host_graph, &shape_inference_graphs,
+ "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
+ host_compute_core, &fld, &shape_inference_graphs,
&has_outside_compilation));
- // Check `host_graph` is empty.
- EXPECT_FALSE(host_graph);
+ // Check host graph is empty.
+ FunctionBody *host_fbody = nullptr;
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> host_func_attrs;
+ host_func_attrs["device_ordinal"] = device_ordinal_temp_value;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &host_fbody));
+ std::unique_ptr<FunctionBody> host_fbody_deleter(host_fbody);
+ Graph *host_graph = host_fbody->graph;
+ EXPECT_EQ(host_graph->num_nodes(), 2);
}
TEST(ExtractOutsideCompilationForFunctionTest, XlaHostComputeRemoved) {
// Build the XLA computation func.
// "const0"
- // "const1" (outside compilation clsuter "0")
+ // "const1" (outside compilation cluster "0")
FunctionDefLibrary fdl;
{
tensorflow::Scope s = tensorflow::Scope::NewRootScope();
@@ -401,31 +427,43 @@
protobuf::Map<string, tensorflow::AttrValue> attrs;
std::map<string, int> host_compute_core = {{"0", 1}, {"1", 0}};
- std::unique_ptr<Graph> host_graph;
std::vector<string> shape_inference_graphs;
bool has_outside_compilation;
NameAttrList name_attrs;
name_attrs.set_name("cluster");
*name_attrs.mutable_attr() = attrs;
TF_CHECK_OK(ExtractOutsideCompilationForFunction(
- "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten",
- host_compute_core, &fld, &host_graph, &shape_inference_graphs,
+ "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
+ host_compute_core, &fld, &shape_inference_graphs,
&has_outside_compilation));
// Check rewritten XLA graph: verify that we have no XlaHostCompute.
- FunctionBody *fbody = nullptr;
- TF_CHECK_OK(FunctionDefToBodyHelper(*fld.Find("cluster_rewritten"),
- AttrSlice(), &fld,
- [&](const string &op, const OpDef **sig) {
- return fld.LookUpOpDef(op, sig);
- },
- &fbody));
- std::unique_ptr<FunctionBody> fbody_deleter(fbody);
- for (Node *n : fbody->graph->nodes()) {
+ FunctionBody *xla_fbody = nullptr;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("cluster_rewritten"), AttrSlice(), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &xla_fbody));
+ std::unique_ptr<FunctionBody> xla_fbody_deleter(xla_fbody);
+ for (Node *n : xla_fbody->graph->nodes()) {
EXPECT_NE(n->type_string(), "XlaHostCompute");
}
- // Check `host_graph`: verify we have no placeholder, but we have "const1".
+ // Check host graph: verify we have no placeholder, but we have "const1".
+ FunctionBody *host_fbody = nullptr;
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> host_func_attrs;
+ host_func_attrs["device_ordinal"] = device_ordinal_temp_value;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &host_fbody));
+ std::unique_ptr<FunctionBody> host_fbody_deleter(host_fbody);
+ Graph *host_graph = host_fbody->graph;
int num_key_placeholders = 0;
for (Node *n : host_graph->nodes()) {
if (n->type_string() == "Placeholder" &&
@@ -438,4 +476,301 @@
EXPECT_NE(node_name_index.find("const1"), node_name_index.end());
}
+REGISTER_OP("XlaSendToHost")
+ .Input("input: Tinput")
+ .Attr("Tinput: type")
+ .Attr("key: string")
+ .SetIsStateful();
+
+REGISTER_OP("XlaRecvFromHost")
+ .Output("output: Toutput")
+ .Attr("Toutput: type")
+ .Attr("shape: shape")
+ .Attr("key: string")
+ .SetIsStateful();
+
+TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) {
+ // Build the XLA computation func.
+ // "const0" (bool)
+ // "const1" (int32)
+ // "if0" (pred = "const0", input = "const1", then_branch = "true_fn",
+ // else_branch = "false_fn")
+ FunctionDefLibrary fdl;
+ {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output arg = ops::_Arg(s.WithOpName("arg"), DT_INT32, 0);
+ Output identity = ops::Identity(s.WithOpName("identity_true_fn"), arg);
+ ops::_Retval retval(s.WithOpName("retval"), identity, 0);
+ std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
+ TF_CHECK_OK(s.ToGraph(g.get()));
+ auto node_name_image = g->BuildNodeNameIndex();
+ node_name_image["identity_true_fn"]->AddAttr("_oc", "0");
+ PartialTensorShape shape({2});
+ node_name_image["identity_true_fn"]->AddAttr(
+ kXlaInferredShapesAttrName, std::vector<PartialTensorShape>{shape});
+
+ FunctionDef *true_fn_fdef = fdl.add_function();
+ TF_CHECK_OK(GraphToFunctionDef(*g, "true_fn", true_fn_fdef));
+ }
+ {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output arg = ops::_Arg(s.WithOpName("arg"), DT_INT32, 0);
+ Output identity = ops::Identity(s.WithOpName("identity_false_fn"), arg);
+ ops::_Retval retval(s.WithOpName("retval"), identity, 0);
+ std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
+ TF_CHECK_OK(s.ToGraph(g.get()));
+ auto node_name_image = g->BuildNodeNameIndex();
+ node_name_image["identity_false_fn"]->AddAttr("_oc", "0");
+ PartialTensorShape shape({2});
+ node_name_image["identity_false_fn"]->AddAttr(
+ kXlaInferredShapesAttrName, std::vector<PartialTensorShape>{shape});
+
+ FunctionDef *false_fn_fdef = fdl.add_function();
+ TF_CHECK_OK(GraphToFunctionDef(*g, "false_fn", false_fn_fdef));
+ }
+ {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output cond = ops::Const(s.WithOpName("const0"), true, {2});
+ Output input = ops::Const(s.WithOpName("const1"), 1, {2});
+ NameAttrList true_fn;
+ true_fn.set_name("true_fn");
+ NameAttrList false_fn;
+ false_fn.set_name("false_fn");
+ auto if_op = ops::If(s.WithOpName("if"), cond,
+ std::initializer_list<Input>{cond, input}, {DT_INT32},
+ true_fn, false_fn);
+ ops::_Retval retval(s.WithOpName("retval"), if_op.output[0], 0);
+ std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
+ TF_CHECK_OK(s.ToGraph(g.get()));
+
+ FunctionDef *xla_fdef = fdl.add_function();
+ TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef));
+ }
+ FunctionLibraryDefinition fld(OpRegistry::Global(), fdl);
+
+ protobuf::Map<string, tensorflow::AttrValue> attrs;
+ std::map<string, int> host_compute_core;
+ std::vector<string> shape_inference_graphs;
+ bool has_outside_compilation;
+ NameAttrList name_attrs;
+ name_attrs.set_name("cluster");
+ *name_attrs.mutable_attr() = attrs;
+ TF_CHECK_OK(ExtractOutsideCompilationForFunction(
+ "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
+ host_compute_core, &fld, &shape_inference_graphs,
+ &has_outside_compilation));
+
+ // Check host graph.
+ {
+ FunctionBody *host_fbody = nullptr;
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> host_func_attrs;
+ host_func_attrs["device_ordinal"] = device_ordinal_temp_value;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &host_fbody));
+ std::unique_ptr<FunctionBody> host_fbody_deleter(host_fbody);
+ Graph *host_graph = host_fbody->graph;
+ auto node_name_index = host_graph->BuildNodeNameIndex();
+
+ // Verify we have XlaRecvAtHost to receive "If" predicate.
+ Node *recv_if_pred_node = node_name_index["recv_oc_if_pred_if"];
+ EXPECT_NE(recv_if_pred_node, nullptr);
+
+ // Verify we have an "If" to choose outside compilation between then_branch
+ // and else_branch, and it has `recv_if_pred_node` as cond input.
+ Node *if_oc_node = node_name_index["oc_if_if"];
+ EXPECT_NE(if_oc_node, nullptr);
+ Node *if_oc_node_cond_input;
+ TF_CHECK_OK(if_oc_node->input_node(0, &if_oc_node_cond_input));
+ EXPECT_EQ(if_oc_node_cond_input, recv_if_pred_node);
+
+ // Check that then_branch outside compilation has node "identity_true_fn".
+ const FunctionDef *true_def = fld.Find("oc_then_branch_host_if_if");
+ EXPECT_NE(true_def, nullptr);
+ bool has_identity_true_fn_node = false;
+ for (const auto &node_def : true_def->node_def()) {
+ if (node_def.name() == "identity_true_fn") {
+ has_identity_true_fn_node = true;
+ break;
+ }
+ }
+ EXPECT_TRUE(has_identity_true_fn_node);
+
+ // Check that else_branch outside compilation has node "identity_false_fn".
+ const FunctionDef *false_def = fld.Find("oc_else_branch_host_if_if");
+ EXPECT_NE(false_def, nullptr);
+ bool has_identity_false_fn_node = false;
+ for (const auto &node_def : false_def->node_def()) {
+ if (node_def.name() == "identity_false_fn") {
+ has_identity_false_fn_node = true;
+ break;
+ }
+ }
+ EXPECT_TRUE(has_identity_false_fn_node);
+ }
+
+ // Check XLA graph.
+ {
+ FunctionBody *xla_fbody = nullptr;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("cluster_rewritten"), AttrSlice(), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &xla_fbody));
+ std::unique_ptr<FunctionBody> xla_fbody_deleter(xla_fbody);
+ Graph *xla_graph = xla_fbody->graph;
+ auto node_name_index = xla_graph->BuildNodeNameIndex();
+
+ // Check that we have XlaSendToHost to send cond predicate to host.
+ Node *send_if_pred_node = node_name_index["send_oc_if_pred_if"];
+ EXPECT_NE(send_if_pred_node, nullptr);
+
+ // Check that the "If" node now has `send_if_pred_node` as attribute
+ // _xla_token_input_nodes.
+ Node *if_node = node_name_index["if"];
+ EXPECT_NE(if_node, nullptr);
+ std::vector<string> token_inputs;
+ TF_CHECK_OK(
+ GetNodeAttr(if_node->def(), "_xla_token_input_nodes", &token_inputs));
+ EXPECT_THAT(token_inputs, ::testing::ElementsAre("send_oc_if_pred_if"));
+ }
+}
+
+TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInWhile) {
+ // Build the XLA computation func.
+ // "const0" (bool)
+ // "while0" (input = "const0", cond = "cond_fn", body = "body_fn")
+ FunctionDefLibrary fdl;
+ {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output arg = ops::_Arg(s.WithOpName("arg"), DT_BOOL, 0);
+ Output identity = ops::Identity(s.WithOpName("identity_cond_fn"), arg);
+ ops::_Retval retval(s.WithOpName("retval"), identity, 0);
+ std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
+ TF_CHECK_OK(s.ToGraph(g.get()));
+ auto node_name_image = g->BuildNodeNameIndex();
+ node_name_image["identity_cond_fn"]->AddAttr("_oc", "0");
+ PartialTensorShape shape({2});
+ node_name_image["identity_cond_fn"]->AddAttr(
+ kXlaInferredShapesAttrName, std::vector<PartialTensorShape>{shape});
+
+ FunctionDef *cond_fn_fdef = fdl.add_function();
+ TF_CHECK_OK(GraphToFunctionDef(*g, "cond_fn", cond_fn_fdef));
+ }
+ {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output arg = ops::_Arg(s.WithOpName("arg"), DT_BOOL, 0);
+ Output identity = ops::Identity(s.WithOpName("identity_body_fn"), arg);
+ ops::_Retval retval(s.WithOpName("retval"), identity, 0);
+ std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
+ TF_CHECK_OK(s.ToGraph(g.get()));
+ auto node_name_image = g->BuildNodeNameIndex();
+ node_name_image["identity_body_fn"]->AddAttr("_oc", "0");
+ PartialTensorShape shape({2});
+ node_name_image["identity_body_fn"]->AddAttr(
+ kXlaInferredShapesAttrName, std::vector<PartialTensorShape>{shape});
+
+ FunctionDef *body_fn_fdef = fdl.add_function();
+ TF_CHECK_OK(GraphToFunctionDef(*g, "body_fn", body_fn_fdef));
+ }
+ {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output input = ops::Const(s.WithOpName("const0"), true, {2});
+ NameAttrList cond_fn;
+ cond_fn.set_name("cond_fn");
+ NameAttrList body_fn;
+ body_fn.set_name("body_fn");
+ auto while_op =
+ ops::While(s.WithOpName("while"), std::initializer_list<Input>{input},
+ cond_fn, body_fn);
+ ops::_Retval retval(s.WithOpName("retval"), while_op.output[0], 0);
+ std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
+ TF_CHECK_OK(s.ToGraph(g.get()));
+
+ FunctionDef *xla_fdef = fdl.add_function();
+ TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef));
+ }
+ FunctionLibraryDefinition fld(OpRegistry::Global(), fdl);
+
+ protobuf::Map<string, tensorflow::AttrValue> attrs;
+ std::map<string, int> host_compute_core;
+ std::vector<string> shape_inference_graphs;
+ bool has_outside_compilation;
+ NameAttrList name_attrs;
+ name_attrs.set_name("cluster");
+ *name_attrs.mutable_attr() = attrs;
+ TF_CHECK_OK(ExtractOutsideCompilationForFunction(
+ "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
+ host_compute_core, &fld, &shape_inference_graphs,
+ &has_outside_compilation));
+
+ // Check host graph.
+ {
+ FunctionBody *host_fbody = nullptr;
+ AttrValue device_ordinal_temp_value;
+ device_ordinal_temp_value.set_i(0);
+ protobuf::Map<string, AttrValue> host_func_attrs;
+ host_func_attrs["device_ordinal"] = device_ordinal_temp_value;
+ TF_CHECK_OK(FunctionDefToBodyHelper(
+ *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld,
+ [&](const string &op, const OpDef **sig) {
+ return fld.LookUpOpDef(op, sig);
+ },
+ &host_fbody));
+ std::unique_ptr<FunctionBody> host_fbody_deleter(host_fbody);
+ Graph *host_graph = host_fbody->graph;
+ auto node_name_index = host_graph->BuildNodeNameIndex();
+
+ // Verify we have an "While" to execute outside compilation.
+ Node *while_oc_node = node_name_index["oc_while_while"];
+ EXPECT_NE(while_oc_node, nullptr);
+
+ // Check that cond outside compilation has node "identity_cond_fn".
+ const FunctionDef *cond_def = fld.Find("oc_cond_host_while_while");
+ EXPECT_NE(cond_def, nullptr);
+ bool has_identity_cond_fn_node = false;
+ for (const auto &node_def : cond_def->node_def()) {
+ if (node_def.name() == "identity_cond_fn") {
+ has_identity_cond_fn_node = true;
+ break;
+ }
+ }
+ EXPECT_TRUE(has_identity_cond_fn_node);
+
+ // Check that body outside compilation has node "identity_body_fn".
+ const FunctionDef *body_def = fld.Find("oc_body_host_while_while");
+ EXPECT_NE(body_def, nullptr);
+ bool has_identity_body_fn_node = false;
+ for (const auto &node_def : body_def->node_def()) {
+ if (node_def.name() == "identity_body_fn") {
+ has_identity_body_fn_node = true;
+ break;
+ }
+ }
+ EXPECT_TRUE(has_identity_body_fn_node);
+ }
+
+ // Check XLA graph.
+ {
+ // Verify that rewritten cond fn has XlaSendToHost to send loop predicate to
+ // host.
+ const FunctionDef *cond_def = fld.Find("cond_fn_oc");
+ EXPECT_NE(cond_def, nullptr);
+ bool has_send_oc_while_cond_node = false;
+ for (const auto &node_def : cond_def->node_def()) {
+ if (node_def.name() == "send_oc_while_cond_while") {
+ has_send_oc_while_cond_node = true;
+ break;
+ }
+ }
+ EXPECT_TRUE(has_send_oc_while_cond_node);
+ }
+}
+
} // namespace tensorflow
diff --git a/tensorflow/compiler/jit/shape_inference.cc b/tensorflow/compiler/jit/shape_inference.cc
index 80c691f..a27e0d9 100644
--- a/tensorflow/compiler/jit/shape_inference.cc
+++ b/tensorflow/compiler/jit/shape_inference.cc
@@ -53,7 +53,15 @@
// shapes, even if no shape function is registered for a node.
Status status = shape_refiner->AddNode(n);
if (!status.ok()) {
- VLOG(1) << "Shape inference failed for node: " << status;
+ VLOG(1) << "Shape inference failed for node " << n->name() << ": "
+ << status;
+ } else {
+ shape_inference::InferenceContext* context = shape_refiner->GetContext(n);
+ for (int i = 0; i < n->num_outputs(); i++) {
+ shape_inference::ShapeHandle handle = context->output(i);
+ VLOG(4) << "Output " << i << " for node " << n->name() << ": "
+ << context->DebugString(handle);
+ }
}
if (n->type_string() == "_Arg") {
diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
index 1fe612d..c7e8d61 100644
--- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
+++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
@@ -142,11 +142,22 @@
TF_RETURN_IF_ERROR(ctx->allocate_temp(
device_tensor.dtype(), device_tensor.shape(), &host_tensor, attrs));
Notification n;
+ Status status;
ctx->op_device_context()->CopyDeviceTensorToCPU(
&device_tensor, "ConstantArgument",
reinterpret_cast<Device*>(ctx->device()), &host_tensor,
- [&](Status status) { n.Notify(); });
+ [&](Status s) {
+ status = s;
+ n.Notify();
+ });
n.WaitForNotification();
+ if (!status.ok()) {
+ LOG(ERROR) << "Copying tensor of shape "
+ << device_tensor.shape().DebugString() << " from "
+ << ctx->device()->name() << "to CPU failed with "
+ << status.ToString();
+ return status;
+ }
constant_arguments[i] = host_tensor;
}
}
@@ -189,6 +200,7 @@
std::map<int, OptionalTensor> variable_args = GetVariables(ctx);
std::vector<XlaCompiler::Argument> args;
+
TF_RETURN_IF_ERROR(XlaComputationLaunchContext::BuildXlaCompilerArguments(
constant_arguments, variable_args, ctx, &args));
diff --git a/tensorflow/compiler/jit/xla_cpu_device.cc b/tensorflow/compiler/jit/xla_cpu_device.cc
index 7df898a..e977064 100644
--- a/tensorflow/compiler/jit/xla_cpu_device.cc
+++ b/tensorflow/compiler/jit/xla_cpu_device.cc
@@ -63,7 +63,19 @@
options.device_ordinal = 0;
options.compilation_device_name = DEVICE_CPU_XLA_JIT;
options.use_multiple_streams = false;
- devices->push_back(absl::make_unique<XlaDevice>(session_options, options));
+ auto device = absl::make_unique<XlaDevice>(session_options, options);
+
+ // Setting GpuDeviceInfo because eager runtime relies on the device
+ // context in tensorflow_gpu_device_info(). Also,
+ // tensorflow_gpu_device_info() == nullptr is used as an IsCPU test.
+ // We need XlaCpuDevice to be treated not as CPU because it allocates
+ // XlaTensors, not regular Tensors.
+ Status status = device->UseGpuDeviceInfo();
+ if (!status.ok()) {
+ errors::AppendToMessage(&status, "while setting up ", DEVICE_GPU_XLA_JIT);
+ return status;
+ }
+ devices->push_back(std::move(device));
return Status::OK();
}
diff --git a/tensorflow/compiler/jit/xla_device_context.cc b/tensorflow/compiler/jit/xla_device_context.cc
index 6e65327..1f3afe8 100644
--- a/tensorflow/compiler/jit/xla_device_context.cc
+++ b/tensorflow/compiler/jit/xla_device_context.cc
@@ -79,6 +79,13 @@
}
}
+void XlaDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,
+ Device* device,
+ Tensor* output_tensor,
+ StatusCallback done) const {
+ done(errors::Unimplemented("XLA->XLA same-device copies not implemented."));
+}
+
void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
Device* device,
Tensor* device_tensor,
diff --git a/tensorflow/compiler/jit/xla_device_context.h b/tensorflow/compiler/jit/xla_device_context.h
index 1e18df1..e45db98 100644
--- a/tensorflow/compiler/jit/xla_device_context.h
+++ b/tensorflow/compiler/jit/xla_device_context.h
@@ -62,6 +62,9 @@
void CopyDeviceTensorToCPU(const Tensor* device_tensor,
absl::string_view tensor_name, Device* device,
Tensor* cpu_tensor, StatusCallback done) override;
+ void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
+ Tensor* output_tensor,
+ StatusCallback done) const override;
xla::LocalClient* client() const { return client_; }
se::Stream* stream() const { return stream_.get(); }
diff --git a/tensorflow/compiler/jit/xla_device_ops.h b/tensorflow/compiler/jit/xla_device_ops.h
index adf0f99..927f983 100644
--- a/tensorflow/compiler/jit/xla_device_ops.h
+++ b/tensorflow/compiler/jit/xla_device_ops.h
@@ -203,6 +203,8 @@
.HostMemory("output") \
.TypeConstraint<ResourceHandle>("T"), \
ArgOp); \
+ REGISTER_KERNEL_BUILDER( \
+ Name(kArgOp).Device(DEVICE).TypeConstraint<Variant>("T"), ArgOp); \
\
REGISTER_KERNEL_BUILDER(Name(kRetOp) \
.Device(DEVICE) \
diff --git a/tensorflow/compiler/jit/xla_gpu_device.cc b/tensorflow/compiler/jit/xla_gpu_device.cc
index 944f732..0191315 100644
--- a/tensorflow/compiler/jit/xla_gpu_device.cc
+++ b/tensorflow/compiler/jit/xla_gpu_device.cc
@@ -16,7 +16,10 @@
// Registers the XLA_GPU device, which is an XlaDevice instantiation that runs
// operators using XLA via the XLA "CUDA" (GPU) backend.
+#include <set>
#include "absl/memory/memory.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_split.h"
#include "tensorflow/compiler/jit/kernels/xla_ops.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_device_ops.h"
@@ -52,8 +55,35 @@
VLOG(1) << "Failed to create XLA_GPU device: " << platform.status();
return Status::OK();
}
-
- for (int i = 0; i < platform.ValueOrDie()->VisibleDeviceCount(); ++i) {
+ string allowed_gpus =
+ session_options.config.gpu_options().visible_device_list();
+ std::set<int> gpu_ids;
+ int num_visible_devices = platform.ValueOrDie()->VisibleDeviceCount();
+ if (allowed_gpus.empty()) {
+ for (int i = 0; i < num_visible_devices; ++i) {
+ gpu_ids.insert(i);
+ }
+ } else {
+ // For loop below is copied from gpu/gpu_device.cc. It validates
+ // the visible_device_list and populates gpu_ids set.
+ const std::vector<string> visible_devices =
+ absl::StrSplit(allowed_gpus, ',');
+ for (const string& platform_gpu_id_str : visible_devices) {
+ int32 platform_gpu_id;
+ if (!absl::SimpleAtoi(platform_gpu_id_str, &platform_gpu_id)) {
+ return errors::InvalidArgument(
+ "Could not parse entry in 'visible_device_list': '",
+ platform_gpu_id_str, "'. visible_device_list = ", allowed_gpus);
+ }
+ if (platform_gpu_id < 0 || platform_gpu_id >= num_visible_devices) {
+ return errors::InvalidArgument(
+ "'visible_device_list' listed an invalid GPU id '", platform_gpu_id,
+ "' but visible device count is ", num_visible_devices);
+ }
+ gpu_ids.insert(platform_gpu_id);
+ }
+ }
+ for (int i : gpu_ids) {
XlaDevice::Options options;
options.platform = platform.ValueOrDie();
options.device_name_prefix = name_prefix;
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index 706811f..554227f 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -205,7 +205,7 @@
allocator_(allocator) {}
~XlaTensorBuffer() override {
- if (data_) {
+ if (data()) {
allocator_->DeallocateRaw(data());
}
}
@@ -229,7 +229,6 @@
}
private:
- void* data_;
size_t expected_size_;
size_t actual_size_;
Allocator* allocator_;
diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index bc3d60b..093b616 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -408,13 +408,6 @@
name = "eager_test",
size = "large",
srcs = ["eager_test.py"],
- disabled_backends = [
- # TODO(b/78199195) Support XLA CPU devices in eager runtime
- "cpu",
- "cpu_ondemand",
- # TODO(b/78468222) Enable GPU backend
- "gpu",
- ],
deps = [
":xla_test",
"//tensorflow/python:array_ops",
diff --git a/tensorflow/compiler/tests/depthwise_conv_op_test.py b/tensorflow/compiler/tests/depthwise_conv_op_test.py
index 174bfa9..6183d3e 100644
--- a/tensorflow/compiler/tests/depthwise_conv_op_test.py
+++ b/tensorflow/compiler/tests/depthwise_conv_op_test.py
@@ -379,8 +379,8 @@
for index, (input_size, filter_size, output_size, stride,
padding) in enumerate(ConfigsToTest()):
print("Testing DepthwiseConv2DFilterGradCompare,", index, "th config:",
- input_size, "*", filter_size, "stride:", stride, "padding:",
- padding)
+ input_size, "*", filter_size, "producing output", output_size,
+ "stride:", stride, "padding:", padding)
self._CompareBackpropFilter(input_size, filter_size, output_size,
stride, padding)
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index fa51a72..a18a4e9 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -1,16 +1,11 @@
+load("//tensorflow:tensorflow.bzl", "tf_copts", "tf_kernel_library")
+
licenses(["notice"]) # Apache 2.0
package(
default_visibility = ["//tensorflow/compiler/tf2xla:internal"],
)
-load("//tensorflow:tensorflow.bzl", "tf_copts")
-load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
-load(
- "//third_party/mkl:build_defs.bzl",
- "if_mkl",
-)
-
tf_kernel_library(
name = "xla_ops",
srcs = [
@@ -122,13 +117,9 @@
"//tensorflow/compiler/tf2xla:common",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla/lib:broadcast",
- "//tensorflow/compiler/tf2xla/lib:cholesky",
- "//tensorflow/compiler/tf2xla/lib:qr",
"//tensorflow/compiler/tf2xla/lib:random",
"//tensorflow/compiler/tf2xla/lib:scatter",
- "//tensorflow/compiler/tf2xla/lib:triangular_solve",
"//tensorflow/compiler/tf2xla/lib:util",
- "//tensorflow/compiler/tf2xla/lib:while_loop",
"//tensorflow/compiler/tf2xla/ops:xla_ops",
"//tensorflow/compiler/xla:array4d",
"//tensorflow/compiler/xla:literal",
@@ -141,12 +132,16 @@
"//tensorflow/compiler/xla/client:xla_builder",
"//tensorflow/compiler/xla/client:xla_computation",
"//tensorflow/compiler/xla/client/lib:arithmetic",
+ "//tensorflow/compiler/xla/client/lib:cholesky",
"//tensorflow/compiler/xla/client/lib:constants",
+ "//tensorflow/compiler/xla/client/lib:loops",
"//tensorflow/compiler/xla/client/lib:math",
"//tensorflow/compiler/xla/client/lib:matrix",
"//tensorflow/compiler/xla/client/lib:pooling",
"//tensorflow/compiler/xla/client/lib:prng",
+ "//tensorflow/compiler/xla/client/lib:qr",
"//tensorflow/compiler/xla/client/lib:sorting",
+ "//tensorflow/compiler/xla/client/lib:triangular_solve",
"//tensorflow/core:framework",
"//tensorflow/core:image_ops_op_lib",
"//tensorflow/core:lib",
diff --git a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc
index 9fcbc86..0ed3044 100644
--- a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc
@@ -13,9 +13,9 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/cholesky.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/client/lib/cholesky.h"
namespace tensorflow {
namespace {
@@ -24,7 +24,7 @@
public:
explicit CholeskyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
void Compile(XlaOpKernelContext* ctx) override {
- ctx->SetOutput(0, Cholesky(ctx->Input(0)));
+ ctx->SetOutput(0, xla::Cholesky(ctx->Input(0)));
}
};
diff --git a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
index 641fefa..399e6e1 100644
--- a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
+++ b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
@@ -392,23 +392,31 @@
builder->GetShape(activations));
TF_ASSIGN_OR_RETURN(xla::Shape out_backprop_shape,
builder->GetShape(gradients));
+ xla::XlaOp filter_backprop;
+
+ xla::Shape input_shape = activations_shape;
+ xla::Shape output_shape = out_backprop_shape;
+
+ TensorShape input_tensor_shape, filter_tensor_shape, output_tensor_shape;
+ TF_RETURN_IF_ERROR(XLAShapeToTensorShape(filter_shape, &filter_tensor_shape));
+ TF_RETURN_IF_ERROR(XLAShapeToTensorShape(input_shape, &input_tensor_shape));
+ TF_RETURN_IF_ERROR(XLAShapeToTensorShape(output_shape, &output_tensor_shape));
+
const xla::Shape expanded_filter_shape =
attrs.depthwise ? ExpandedFilterShapeForDepthwiseConvolution(filter_shape)
: filter_shape;
-
// Reuse dimension computation logic from conv_grad_ops.cc.
ConvBackpropDimensions dims;
+ // The filter gradients are computed by a convolution of the input
+ // activations and the output gradients, with some appropriate padding.
+ // See the comment at the top of conv_grad_ops.h for details.
+ xla::ConvolutionDimensionNumbers dnums;
+
TF_RETURN_IF_ERROR(ConvBackpropComputeDimensionsV2XlaShapes(
type_string, attrs.num_spatial_dims, activations_shape,
expanded_filter_shape, out_backprop_shape, attrs.dilations, attrs.strides,
attrs.padding, attrs.data_format, &dims));
- // The filter gradients are computed by a convolution of the input
- // activations and the output gradients, with some appropriate padding.
- // See the comment at the top of conv_grad_ops.h for details.
-
- xla::ConvolutionDimensionNumbers dnums;
-
// The activations (inputs) form the LHS of the convolution.
// Activations have shape: [batch, in_rows, in_cols, ..., in_depth]
// For the gradient computation, we flip the roles of the batch and
@@ -420,29 +428,97 @@
int n_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format);
int c_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format);
- // Swap n_dim and c_dim in the activations.
- dnums.set_input_batch_dimension(c_dim);
- dnums.set_input_feature_dimension(n_dim);
+ int64 total_spatial_size = 1;
+ for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+ total_spatial_size *= dims.input_size(i);
+ }
- // The gradients become the RHS of the convolution.
- // The gradients have shape [batch, out_rows, out_cols, ..., out_depth]
- // where the batch becomes the input feature for the convolution.
- dnums.set_kernel_input_feature_dimension(n_dim);
- dnums.set_kernel_output_feature_dimension(c_dim);
+ // We use this approach only for depthwise convolutions where feature counts
+ // are large but space dimensions are small.
+ bool should_perform_depthwise_conv =
+ (total_spatial_size < dims.in_depth) &&
+ filter_tensor_shape.dim_size(num_dims - 1) == 1 && attrs.depthwise;
- std::vector<std::pair<int64, int64>> padding(attrs.num_spatial_dims);
- std::vector<int64> rhs_dilation(attrs.num_spatial_dims);
- std::vector<int64> window_strides(attrs.num_spatial_dims);
- std::vector<int64> ones(attrs.num_spatial_dims, 1);
+ int64 num_spatial_dims =
+ attrs.num_spatial_dims + (should_perform_depthwise_conv ? 1 : 0);
+
+ std::vector<std::pair<int64, int64>> padding(num_spatial_dims);
+ std::vector<int64> rhs_dilation(num_spatial_dims);
+ std::vector<int64> window_strides(num_spatial_dims);
+ std::vector<int64> ones(num_spatial_dims, 1);
+
+ if (should_perform_depthwise_conv) {
+ // This approach is similar to handling of grouped convolutions in
+ // the convolution_feature_group_converter.cc. Please refer to it for
+ // details.
+
+ // Add spatial dimension to the activation, and reshape.
+ std::vector<int64> activations_reshape_sizes, gradients_reshape_sizes;
+
+ activations_reshape_sizes.push_back(dims.batch_size);
+ gradients_reshape_sizes.push_back(dims.batch_size);
+ for (int i = 0; i < attrs.num_spatial_dims; i++) {
+ activations_reshape_sizes.push_back(dims.input_size(i));
+ gradients_reshape_sizes.push_back(dims.output_size(i));
+ }
+ activations_reshape_sizes.push_back(dims.in_depth);
+ activations_reshape_sizes.push_back(1);
+ gradients_reshape_sizes.push_back(dims.out_depth);
+ gradients_reshape_sizes.push_back(1);
+
+ activations = xla::Reshape(activations, activations_reshape_sizes);
+ gradients = xla::Reshape(gradients, gradients_reshape_sizes);
+
+ int64 new_spatial_dim = activations_reshape_sizes.size() - 1;
+
+ // Set the newly added dimension to be the batch.
+ dnums.set_input_batch_dimension(new_spatial_dim);
+ dnums.set_input_feature_dimension(c_dim);
+
+ // The gradients become the RHS of the convolution.
+ // The gradients have shape [batch, out_rows, out_cols, ..., out_depth, 1]
+ // where the batch becomes a spatial dimension, and 1 becomes
+ // the input feature for the convolution.
+ dnums.set_kernel_input_feature_dimension(new_spatial_dim);
+ dnums.set_kernel_output_feature_dimension(c_dim);
+
+ // Treat original batch dimension as a spatial dimension.
+ dnums.add_input_spatial_dimensions(n_dim);
+ dnums.add_kernel_spatial_dimensions(n_dim);
+ } else {
+ // The activations (inputs) form the LHS of the convolution.
+ // Activations have shape: [batch, in_rows, in_cols, ..., in_depth]
+ // For the gradient computation, we flip the roles of the batch and
+ // feature dimensions.
+ // Each spatial entry has size in_depth * batch
+
+ // Swap n_dim and c_dim in the activations.
+ dnums.set_input_batch_dimension(c_dim);
+ dnums.set_input_feature_dimension(n_dim);
+
+ // The gradients become the RHS of the convolution.
+ // The gradients have shape [batch, out_rows, out_cols, ..., out_depth]
+ // where the batch becomes the input feature for the convolution.
+ dnums.set_kernel_input_feature_dimension(n_dim);
+ dnums.set_kernel_output_feature_dimension(c_dim);
+ }
+
+ dnums.set_output_batch_dimension(num_spatial_dims);
+ dnums.set_output_feature_dimension(num_spatial_dims + 1);
// Tensorflow filter shape is [ H, W, ..., inC, outC ].
- for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+ for (int i = 0; i < num_spatial_dims; ++i) {
dnums.add_output_spatial_dimensions(i);
}
- dnums.set_output_batch_dimension(attrs.num_spatial_dims);
- dnums.set_output_feature_dimension(attrs.num_spatial_dims + 1);
- for (int i = 0; i < attrs.num_spatial_dims; ++i) {
+ if (should_perform_depthwise_conv) {
+ // Set the right parameters for the newly created spatial dimension.
+ padding[0] = {0, 0};
+ rhs_dilation[0] = 1;
+ window_strides[0] = 1;
+ }
+
+ for (int64 i = 0; i < attrs.num_spatial_dims; ++i) {
int64 dim = GetTensorSpatialDimIndex(num_dims, attrs.data_format, i);
dnums.add_input_spatial_dimensions(dim);
dnums.add_kernel_spatial_dimensions(dim);
@@ -483,9 +559,10 @@
const int64 pad_before =
attrs.padding == Padding::SAME ? std::max<int64>(pad_total / 2, 0) : 0;
- padding[i] = {pad_before, pad_total - pad_before};
- rhs_dilation[i] = dims.spatial_dims[i].stride;
- window_strides[i] = attrs.dilations[dim];
+ int64 dim_being_operated = should_perform_depthwise_conv ? i + 1 : i;
+ padding[dim_being_operated] = {pad_before, pad_total - pad_before};
+ rhs_dilation[dim_being_operated] = dims.spatial_dims[i].stride;
+ window_strides[dim_being_operated] = attrs.dilations[dim];
}
// Besides padding the input, we will also expand output_rows to
@@ -496,13 +573,19 @@
//
// This is done by specifying the window dilation factors in the
// convolution HLO below.
- auto filter_backprop =
- xla::ConvGeneralDilated(activations, gradients, window_strides, padding,
- /*lhs_dilation=*/ones, rhs_dilation, dnums);
+ filter_backprop = xla::ConvGeneralDilated(
+ activations, gradients, window_strides, padding,
+ /*lhs_dilation=*/ones, rhs_dilation, dnums,
+ /*feature_group_count=*/
+ should_perform_depthwise_conv ? dims.in_depth : 1);
- if (attrs.depthwise) {
- filter_backprop = ContractFilterForDepthwiseBackprop(
- filter_shape, filter_backprop, activations.builder());
+ if (should_perform_depthwise_conv) {
+ filter_backprop = xla::Reshape(filter_backprop, filter_shape.dimensions());
+ } else {
+ if (attrs.depthwise) {
+ filter_backprop = ContractFilterForDepthwiseBackprop(
+ filter_shape, filter_backprop, activations.builder());
+ }
}
return filter_backprop;
diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc
index 20b0de1..41c31d0 100644
--- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc
@@ -14,7 +14,6 @@
==============================================================================*/
#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h"
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/tf2xla/type_util.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index e9bb0a7..96ddd42 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -15,12 +15,12 @@
#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h"
#include "tensorflow/compiler/tf2xla/lib/util.h"
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
+#include "tensorflow/compiler/xla/client/lib/loops.h"
#include "tensorflow/compiler/xla/client/lib/sorting.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/shape_util.h"
@@ -505,9 +505,9 @@
init_values.push_back(included_iou);
auto suppress_loop_result =
- XlaWhileLoop(WhileCondFn(num_boxes, output_size),
- SuppressBodyFn(num_boxes), init_values, "suppress_loop",
- builder)
+ xla::WhileLoopHelper(WhileCondFn(num_boxes, output_size),
+ SuppressBodyFn(num_boxes), init_values,
+ "suppress_loop", builder)
.ValueOrDie();
xla::XlaOp included_score =
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
index f4def11..90c0ebe 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
@@ -13,9 +13,9 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/client/lib/triangular_solve.h"
namespace tensorflow {
namespace {
@@ -29,7 +29,7 @@
}
void Compile(XlaOpKernelContext* ctx) override {
- auto result = TriangularSolve(
+ auto result = xla::TriangularSolve(
ctx->Input(0), ctx->Input(1), /*left_side=*/true,
/*lower=*/lower_, /*transpose_a=*/adjoint_, /*conjugate_a=*/adjoint_);
ctx->SetOutput(0, result);
diff --git a/tensorflow/compiler/tf2xla/kernels/qr_op.cc b/tensorflow/compiler/tf2xla/kernels/qr_op.cc
index 7ea0afc..66ec40a 100644
--- a/tensorflow/compiler/tf2xla/kernels/qr_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/qr_op.cc
@@ -13,9 +13,9 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/qr.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/client/lib/qr.h"
namespace tensorflow {
namespace {
@@ -26,7 +26,7 @@
OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_));
}
void Compile(XlaOpKernelContext* ctx) override {
- auto result = QRDecomposition(ctx->Input(0), full_matrices_);
+ auto result = xla::QRDecomposition(ctx->Input(0), full_matrices_);
if (!result.ok()) {
ctx->SetStatus(result.status());
return;
diff --git a/tensorflow/compiler/tf2xla/kernels/random_ops.cc b/tensorflow/compiler/tf2xla/kernels/random_ops.cc
index 8822e29..2d92056 100644
--- a/tensorflow/compiler/tf2xla/kernels/random_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/random_ops.cc
@@ -20,12 +20,12 @@
#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h"
#include "tensorflow/compiler/tf2xla/lib/random.h"
#include "tensorflow/compiler/tf2xla/lib/util.h"
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
+#include "tensorflow/compiler/xla/client/lib/loops.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
@@ -175,8 +175,8 @@
};
// for i in range(n):
auto swap_loop_result =
- XlaForEachIndex(n, xla::S32, swap_body_fn, {swaps, indices},
- "indices_swap_loop", builder)
+ xla::ForEachIndex(n, xla::S32, swap_body_fn, {swaps, indices},
+ "indices_swap_loop", builder)
.ValueOrDie();
auto swapped_indices = swap_loop_result[1];
diff --git a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc
index 54d34a3..f9985d5 100644
--- a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc
@@ -125,7 +125,7 @@
dimensions.back() = 1;
auto batch_indices =
- xla::Iota(b, xla::ShapeUtil::MakeShape(xla::U32, dimensions),
+ xla::Iota(b, xla::ShapeUtil::MakeShape(xla::S32, dimensions),
/*iota_dimension=*/0);
return xla::ConcatInDim(b, {batch_indices, indices}, dimensions.size() - 1);
@@ -189,11 +189,53 @@
scatter_dim_numbers);
}
+// Bounds samples to 0 if the warp image indices are out of the (-1, image_size)
+// bound.
+// The resulting dimension is given by 'result_dims'.
+XlaOp BoundSamples(XlaOpKernelContext* ctx, XlaOp warp,
+ xla::PrimitiveType warp_type, TensorShape warp_shape,
+ std::vector<int64> result_dims,
+ std::vector<int64> broadcasted_dims, int64 last_warp_dim,
+ xla::Shape data_shape, XlaOp sample) {
+ auto is_gt_minus_one =
+ xla::Gt(warp,
+ xla::ConvertElementType(
+ xla::ConstantR1<float>(ctx->builder(), {-1, -1}), warp_type),
+ /*broadcast_dimensions=*/{warp_shape.dims() - 1});
+ auto is_lt_image_size = xla::Lt(
+ warp,
+ xla::ConvertElementType(
+ xla::ConstantR1<float>(
+ ctx->builder(),
+ {/*width=*/static_cast<float>(data_shape.dimensions(2)),
+ /*height=*/static_cast<float>(data_shape.dimensions(1))}),
+ warp_type),
+ /*broadcast_dimensions=*/{warp_shape.dims() - 1});
+
+ auto is_in_bound_padded_x_y = xla::And(is_gt_minus_one, is_lt_image_size);
+ // Reduce along last dimension. The resulting dimension is:
+ // [batch, dim_0, ...dim_n].
+ auto is_in_bound = xla::Reduce(
+ is_in_bound_padded_x_y, xla::ConstantR0<bool>(ctx->builder(), true),
+ xla::CreateScalarAndComputation(xla::PrimitiveType::PRED, ctx->builder()),
+ {last_warp_dim});
+
+ // Broadcast 'is_in_bound' to the same dimension as 'result_dims'.
+ auto broadcasted_is_in_bound =
+ xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims);
+
+ // Set out of bound samples to zero.
+ auto zeros =
+ xla::Broadcast(xla::Zero(ctx->builder(), warp_type), result_dims);
+ return xla::Select(broadcasted_is_in_bound, sample, zeros);
+}
+
// Build computation the backprop into input 'data'.
// Where input:
// grad_output is of dimension [batch, dim_0, ...dim_n, channel]
// ratio is of dimension [batch, dim_0, ...dim_n, 2]
// gather_indices is of dimension [batch, dim_0, ...dim_n, 3]
+// data_shape is of dimension [batch, x(width), y(height), channel]
//
// Output:
// scatter-add to each 2x2 grad_data neighbor:
@@ -201,10 +243,12 @@
// grad_data[cx, fy, chan] += output_grad * (1 - dx) * dy
// grad_data[fx, cy, chan] += output_grad * dx * (1 - dy)
// grad_data[cx, cy, chan] += output_grad * (1 - dx) * (1 - dy)
-// where (dx, dy) is (1 - ratio).
+// where (dx, dy) is (1 - ratio). If (dx, dy) is out of bound, then the their
+// contribution is 0 to 'grad_data'.
XlaOp CalculateGradData(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio,
- XlaOp gather_indices, xla::PrimitiveType warp_type,
- TensorShape warp_shape, int64 data_channels,
+ XlaOp gather_indices, XlaOp warp,
+ xla::PrimitiveType warp_type, TensorShape warp_shape,
+ int64 last_warp_dim, int64 data_channels,
xla::Shape data_shape) {
// Weights tensor has dimension [batch, dim_0, ... dim_n, 4].
auto weights = BilinearWeights(ctx, ratio, warp_shape, warp_type);
@@ -229,6 +273,18 @@
std::iota(reshaped_weights_indices.begin(), reshaped_weights_indices.end(),
0);
+ // Set out of bound weights to 0.
+ // The dimension of the reshaped_weight: [batch, dim_0, ...dim_n, 2, 2].
+ std::vector<int64> reshaped_result_dims(warp_dims.begin(),
+ warp_dims.end() - 1);
+ reshaped_result_dims.push_back(2);
+ reshaped_result_dims.push_back(2);
+ std::vector<int64> broadcasted_dims(warp_dims.size() - 1);
+ std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
+ reshaped_weights = BoundSamples(ctx, warp, warp_type, warp_shape,
+ reshaped_result_dims, broadcasted_dims,
+ last_warp_dim, data_shape, reshaped_weights);
+
// The dimension is [batch, dim_0, ..., dim_n, 2, 2, data_channel].
auto broadcast_reshaped_weights = xla::BroadcastInDim(
reshaped_weights, weights_with_channels_dims, reshaped_weights_indices);
@@ -245,18 +301,41 @@
auto grad_data = xla::ConstantLiteral(
ctx->builder(), xla::Literal::CreateFromShape(data_shape));
- return ScatterToGradData(ctx, grad_data, gather_indices,
- grad_output_multiply_weights, warp_shape.dims(),
- warp_type);
+ // Pad grad data then slice it back.
+ //
+ // After left and right column 0-padding, the new dimension of padded data
+ // will be [batch, x+2, y+2, channel].
+ auto padded_grad_data =
+ xla::Pad(grad_data, xla::Zero(ctx->builder(), warp_type),
+ xla::MakeEdgePaddingConfig({{0, 0}, {1, 1}, {1, 1}, {0, 0}}));
+
+ auto shifting_value = xla::ConstantR1<int32>(
+ ctx->builder(), {/*batch=*/0, /*x(width)=*/1, /*y(height)=*/1});
+ auto shifted_gather_indices =
+ xla::Add(gather_indices, shifting_value, {last_warp_dim});
+
+ auto updated_grad_data = ScatterToGradData(
+ ctx, padded_grad_data, shifted_gather_indices,
+ grad_output_multiply_weights, warp_shape.dims(), warp_type);
+
+ const int64 batch_size = data_shape.dimensions(0);
+ const int64 width = data_shape.dimensions(1);
+ const int64 height = data_shape.dimensions(2);
+ // Slice out the result accounting for the padding.
+ return xla::Slice(
+ updated_grad_data, /*start_indices=*/{0, 1, 1, 0},
+ /*limit_indices=*/{batch_size, width + 1, height + 1, data_channels},
+ /*strides=*/{1, 1, 1, 1});
}
// Build computation for the backprop into input 'warp'.
// Where input:
-// warp is of dimension [batch, dim_0, ...dim_n, 2]
-// grad_output is of dimension [batch, dim_0, ...dim_n, channel]
-// ratio is of dimension [batch, dim_0, ...dim_n, 2]
-// gather_indices is of dimension [batch, dim_0, ...dim_n, 3]
-// data is of dimension [batch, x, y, channel]
+// warp is of dimension [batch, dim_0, ...dim_n, 2]
+// grad_output is of dimension [batch, dim_0, ...dim_n, channel]
+// ratio is of dimension [batch, dim_0, ...dim_n, 2]
+// gather_indices is of dimension [batch, dim_0, ...dim_n, 3] where the last
+// dimension of size 3 is for {batch, x(width), y(height)}.
+// data is of dimension [batch, x, y, channel]
//
// Output (simplified by ignoring the batch dimensions):
// Since the forward path has:
@@ -275,12 +354,12 @@
// grad_warp_x = py * (img_cxcy - img_fxcy) + (1-py) * (img_cxfy-img_fxfy)
// grad_warp_y = px * (img_cxcy - img_cxfy) + (1-px) * (img_fxcy-img_fxfy)
//
-// where (px, py) is warp, (fx, fy) is the left top corner and (cx, cy) is the
+// where (px, py) is warp, (fx, fy) is the top left corner and (cx, cy) is the
// bottom right corner in a 2x2 neighborhood.
XlaOp CalculateGradWarp(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio,
XlaOp gather_indices, XlaOp data,
TensorShape warp_shape, int64 data_channels,
- xla::PrimitiveType data_type) {
+ xla::PrimitiveType data_type, xla::Shape data_shape) {
auto warp_dims = warp_shape.dim_sizes();
std::vector<int64> warp_dims_without_last_dims(warp_dims.begin(),
warp_dims.end() - 1);
@@ -289,12 +368,30 @@
std::vector<int64> neighbor_broadcast_dims = warp_dims_without_last_dims;
neighbor_broadcast_dims.push_back(4);
- // The dimension is [batch, dim_0, ... dim_n, 4, data_channels]
- auto neighbors_data = Gather2by2Neighbors(
- ctx->builder(), data, gather_indices, data_channels, warp_shape.dims());
+ // With dimension [batch, dim_0, ...dim_n, 4]
+ auto neighbor_broadcast_shape =
+ xla::ShapeUtil::MakeShape(data_type, neighbor_broadcast_dims);
const int64 last_warp_dim = warp_shape.dims() - 1;
+ // Pad data with 0, before gathering such that 0 will be returned for samples
+ // in the range of (-1, 0) or (image_dimension-1, image_dimension).
+ // After left and right column 0-padding, the new dimension of padded data
+ // will be [batch, x+2, y+2, channel].
+ auto padded_data =
+ xla::Pad(data, xla::Zero(ctx->builder(), data_type),
+ xla::MakeEdgePaddingConfig({{0, 0}, {1, 1}, {1, 1}, {0, 0}}));
+
+ auto shifting_value = xla::ConstantR1<int32>(
+ ctx->builder(), {/*batch=*/0, /*x(width)=*/1, /*y(height)=*/1});
+ auto shifted_gather_indices =
+ xla::Add(gather_indices, shifting_value, {last_warp_dim});
+
+ // The dimension is [batch, dim_0, ... dim_n, 4, data_channels]
+ auto neighbors_data =
+ Gather2by2Neighbors(ctx->builder(), padded_data, shifted_gather_indices,
+ data_channels, warp_shape.dims());
+
// Since we will be creating the dot product of:
// lhs: [batch, dim_0, ...dim_n, 4]
// and
@@ -417,7 +514,7 @@
// Find the coordinates of the top left corner for the 2x2 region to be
// sampled from. The dimensions are [batch, dim_0, ... dim_n, 2] where the
// last dimension of size 2 in turn is [x, y].
- XlaOp top_left = xla::ConvertElementType(warp, xla::U32);
+ XlaOp top_left = xla::ConvertElementType(warp, xla::S32);
auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape);
@@ -526,7 +623,8 @@
size, "]"));
}
// Last dimension of warp shape must be of size 2.
- OP_REQUIRES(ctx, warp_shape.dim_size(warp_shape.dims() - 1) == 2,
+ const int64 last_warp_dim = warp_shape.dims() - 1;
+ OP_REQUIRES(ctx, warp_shape.dim_size(last_warp_dim) == 2,
errors::InvalidArgument(
"the last dimension of warp must be exactly size 2."));
xla::PrimitiveType warp_type = ctx->input_xla_type(1);
@@ -549,24 +647,32 @@
// Find the top left corner coordinate for the region to be sampled from.
// The dimensions are [batch, dim_0, ... dim_n, 2] where the last dimension
// of size 2 in turn is [x, y].
- XlaOp top_left = xla::ConvertElementType(warp, xla::U32);
+ XlaOp top_left = xla::ConvertElementType(xla::Floor(warp), xla::S32);
- // Dimensions are [batch, dim_0, ... dim_n, 2]
+ // Dimensions are [batch, dim_0, ... dim_n, 2].
XlaOp ratio = warp - xla::ConvertElementType(top_left, warp_type);
// Indices for gathering neighboring pixels.
auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape);
- auto grad_data =
- CalculateGradData(ctx, grad_output, ratio, gather_indices, warp_type,
- warp_shape, data_channels, data_shape);
+ auto grad_data = CalculateGradData(
+ ctx, grad_output, ratio, gather_indices, warp, warp_type, warp_shape,
+ last_warp_dim, data_channels, data_shape);
auto grad_warp =
CalculateGradWarp(ctx, grad_output, ratio, gather_indices, data,
- warp_shape, data_channels, data_type);
+ warp_shape, data_channels, data_type, data_shape);
+ auto warp_dims = warp_shape.dim_sizes();
+ std::vector<int64> result_dims(warp_dims.begin(), warp_dims.end() - 1);
+ result_dims.push_back(2);
+ std::vector<int64> broadcasted_dims(warp_dims.size() - 1);
+ std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
+ auto grad_warp_bounded =
+ BoundSamples(ctx, warp, warp_type, warp_shape, result_dims,
+ broadcasted_dims, last_warp_dim, data_shape, grad_warp);
ctx->SetOutput(0, grad_data);
- ctx->SetOutput(1, grad_warp);
+ ctx->SetOutput(1, grad_warp_bounded);
}
};
diff --git a/tensorflow/compiler/tf2xla/kernels/training_ops.cc b/tensorflow/compiler/tf2xla/kernels/training_ops.cc
index 960c146..26d4214 100644
--- a/tensorflow/compiler/tf2xla/kernels/training_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/training_ops.cc
@@ -172,6 +172,65 @@
REGISTER_XLA_OP(Name("ResourceApplyMomentum").TypeConstraint("T", kFloatTypes),
ResourceApplyMomentum);
+class ResourceApplyKerasMomentum : public XlaOpKernel {
+ public:
+ explicit ResourceApplyKerasMomentum(OpKernelConstruction* ctx)
+ : XlaOpKernel(ctx) {
+ OP_REQUIRES_OK(ctx, ctx->GetAttr("use_nesterov", &use_nesterov_));
+ }
+
+ void Compile(XlaOpKernelContext* ctx) override {
+ DataType type = ctx->input_type(2);
+
+ TensorShape var_shape, accum_shape;
+ xla::XlaOp var, accum;
+ OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var));
+ OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &accum_shape, &accum));
+
+ OP_REQUIRES(ctx, var_shape.IsSameSize(accum_shape),
+ errors::InvalidArgument(
+ "var and accum do not have the same shape",
+ var_shape.DebugString(), " ", accum_shape.DebugString()));
+
+ TensorShape lr_shape = ctx->InputShape(2);
+ OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_shape),
+ errors::InvalidArgument("lr is not a scalar: ",
+ lr_shape.DebugString()));
+
+ TensorShape grad_shape = ctx->InputShape(3);
+ OP_REQUIRES(ctx, var_shape.IsSameSize(grad_shape),
+ errors::InvalidArgument(
+ "var and grad do not have the same shape",
+ var_shape.DebugString(), " ", grad_shape.DebugString()));
+
+ TensorShape momentum_shape = ctx->InputShape(4);
+ OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(momentum_shape),
+ errors::InvalidArgument("momentum is not a scalar: ",
+ momentum_shape.DebugString()));
+
+ xla::XlaOp lr = ctx->Input(2);
+ xla::XlaOp grad = ctx->Input(3);
+ xla::XlaOp momentum = ctx->Input(4);
+
+ accum = accum * momentum - grad * lr;
+ if (use_nesterov_) {
+ // See https://github.com/tensorflow/tensorflow/pull/2798 for an
+ // explanation of the reparameterization used here.
+ var = var + accum * momentum - grad * lr;
+ } else {
+ var = var + accum;
+ }
+ OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, var));
+ OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, type, accum));
+ }
+
+ private:
+ bool use_nesterov_;
+};
+REGISTER_XLA_OP(
+ Name("ResourceApplyKerasMomentum").TypeConstraint("T", kFloatTypes),
+ ResourceApplyKerasMomentum);
+
class ResourceApplyAdagrad : public XlaOpKernel {
public:
explicit ResourceApplyAdagrad(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc
index ce007fc..89b577b 100644
--- a/tensorflow/compiler/tf2xla/kernels/while_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc
@@ -41,8 +41,7 @@
*has_uninitialized_vars = false;
*has_tensor_arrays = false;
for (int i = 0; i < ctx->num_inputs(); ++i) {
- VLOG(2) << " Input " << i
- << " type: " << DataTypeString(ctx->input_type(i))
+ VLOG(2) << " Input " << i << " type: " << DataTypeString(ctx->input_type(i))
<< " shape: " << ctx->InputShape(i).DebugString();
XlaCompiler::Argument& arg = (*args)[i];
DataType type = ctx->input_type(i);
@@ -233,13 +232,22 @@
xla::ShapeUtil::HumanString(body_input_shape), " vs. ",
xla::ShapeUtil::HumanString(body.xla_output_shape)));
- xla::Shape expected_cond_output_shape = xla::ShapeUtil::MakeTupleShape(
- {xla::ShapeUtil::MakeShape(xla::PRED, {})});
+ xla::Shape expected_cond_output_shape_without_side_effect =
+ xla::ShapeUtil::MakeTupleShape(
+ {xla::ShapeUtil::MakeShape(xla::PRED, {})});
+ xla::Shape expected_cond_output_shape_with_side_effect =
+ xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::PRED, {}),
+ xla::ShapeUtil::MakeTokenShape()});
OP_REQUIRES(ctx,
- xla::ShapeUtil::Compatible(cond.xla_output_shape,
- expected_cond_output_shape),
+ xla::ShapeUtil::Compatible(
+ cond.xla_output_shape,
+ expected_cond_output_shape_without_side_effect) ||
+ xla::ShapeUtil::Compatible(
+ cond.xla_output_shape,
+ expected_cond_output_shape_with_side_effect),
errors::InvalidArgument(
- "Output shape of loop condition should be (pred[]), got: ",
+ "Output shape of loop condition should be (pred[]) or "
+ "(pred[], token[]), got: ",
xla::ShapeUtil::HumanString(cond.xla_output_shape)));
int num_inputs = body.input_mapping.size();
diff --git a/tensorflow/compiler/tf2xla/lib/BUILD b/tensorflow/compiler/tf2xla/lib/BUILD
index 422781d..3d7b0bc 100644
--- a/tensorflow/compiler/tf2xla/lib/BUILD
+++ b/tensorflow/compiler/tf2xla/lib/BUILD
@@ -15,8 +15,6 @@
]),
)
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test")
-
cc_library(
name = "broadcast",
srcs = ["broadcast.cc"],
@@ -34,27 +32,6 @@
)
cc_library(
- name = "cholesky",
- srcs = ["cholesky.cc"],
- hdrs = ["cholesky.h"],
- deps = [
- ":triangular_solve",
- ":util",
- ":while_loop",
- "//tensorflow/compiler/xla:literal",
- "//tensorflow/compiler/xla:shape_util",
- "//tensorflow/compiler/xla:status_macros",
- "//tensorflow/compiler/xla:statusor",
- "//tensorflow/compiler/xla:xla_data_proto",
- "//tensorflow/compiler/xla/client:xla_builder",
- "//tensorflow/compiler/xla/client/lib:constants",
- "//tensorflow/compiler/xla/client/lib:matrix",
- "//tensorflow/compiler/xla/client/lib:slicing",
- "//tensorflow/core:lib",
- ],
-)
-
-cc_library(
name = "random",
srcs = ["random.cc"],
hdrs = ["random.h"],
@@ -70,34 +47,11 @@
)
cc_library(
- name = "qr",
- srcs = ["qr.cc"],
- hdrs = ["qr.h"],
- deps = [
- ":util",
- ":while_loop",
- "//tensorflow/compiler/xla:literal_util",
- "//tensorflow/compiler/xla:shape_util",
- "//tensorflow/compiler/xla:status_macros",
- "//tensorflow/compiler/xla:statusor",
- "//tensorflow/compiler/xla:xla_data_proto",
- "//tensorflow/compiler/xla/client:xla_builder",
- "//tensorflow/compiler/xla/client/lib:arithmetic",
- "//tensorflow/compiler/xla/client/lib:constants",
- "//tensorflow/compiler/xla/client/lib:math",
- "//tensorflow/compiler/xla/client/lib:matrix",
- "//tensorflow/compiler/xla/client/lib:slicing",
- "//tensorflow/core:lib",
- ],
-)
-
-cc_library(
name = "scatter",
srcs = ["scatter.cc"],
hdrs = ["scatter.h"],
deps = [
":util",
- ":while_loop",
"//tensorflow/compiler/xla:literal",
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/compiler/xla:status_macros",
@@ -112,52 +66,6 @@
)
cc_library(
- name = "triangular_solve",
- srcs = ["triangular_solve.cc"],
- hdrs = ["triangular_solve.h"],
- deps = [
- ":util",
- "//tensorflow/compiler/xla:literal",
- "//tensorflow/compiler/xla:shape_util",
- "//tensorflow/compiler/xla:status_macros",
- "//tensorflow/compiler/xla:statusor",
- "//tensorflow/compiler/xla:util",
- "//tensorflow/compiler/xla:xla_data_proto",
- "//tensorflow/compiler/xla/client:xla_builder",
- "//tensorflow/compiler/xla/client:xla_computation",
- "//tensorflow/compiler/xla/client/lib:constants",
- "//tensorflow/compiler/xla/client/lib:math",
- "//tensorflow/compiler/xla/client/lib:matrix",
- "//tensorflow/compiler/xla/client/lib:slicing",
- "//tensorflow/core:lib",
- ],
-)
-
-xla_test(
- name = "triangular_solve_test",
- srcs = ["triangular_solve_test.cc"],
- tags = ["noasan"], # sometimes times out, http://b/78650012
- deps = [
- ":triangular_solve",
- "//tensorflow/compiler/xla:array2d",
- "//tensorflow/compiler/xla:literal",
- "//tensorflow/compiler/xla:shape_util",
- "//tensorflow/compiler/xla:statusor",
- "//tensorflow/compiler/xla:test",
- "//tensorflow/compiler/xla:types",
- "//tensorflow/compiler/xla:xla_data_proto",
- "//tensorflow/compiler/xla/client:global_data",
- "//tensorflow/compiler/xla/client:local_client",
- "//tensorflow/compiler/xla/client:xla_builder",
- "//tensorflow/compiler/xla/tests:client_library_test_base",
- "//tensorflow/compiler/xla/tests:literal_test_util",
- "//tensorflow/compiler/xla/tests:xla_internal_test_main",
- "//tensorflow/core:lib",
- "//tensorflow/core:test",
- ],
-)
-
-cc_library(
name = "util",
srcs = ["util.cc"],
hdrs = ["util.h"],
@@ -174,19 +82,3 @@
"@com_google_absl//absl/types:span",
],
)
-
-cc_library(
- name = "while_loop",
- srcs = ["while_loop.cc"],
- hdrs = ["while_loop.h"],
- deps = [
- ":util",
- "//tensorflow/compiler/xla:shape_util",
- "//tensorflow/compiler/xla:status_macros",
- "//tensorflow/compiler/xla:statusor",
- "//tensorflow/compiler/xla/client:xla_builder",
- "//tensorflow/compiler/xla/client:xla_computation",
- "@com_google_absl//absl/strings",
- "@com_google_absl//absl/types:span",
- ],
-)
diff --git a/tensorflow/compiler/tf2xla/lib/scatter.cc b/tensorflow/compiler/tf2xla/lib/scatter.cc
index 2b1c2ce..6880567 100644
--- a/tensorflow/compiler/tf2xla/lib/scatter.cc
+++ b/tensorflow/compiler/tf2xla/lib/scatter.cc
@@ -20,7 +20,6 @@
#include "absl/types/span.h"
#include "tensorflow/compiler/tf2xla/lib/util.h"
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/literal.h"
diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.cc b/tensorflow/compiler/tf2xla/lib/while_loop.cc
deleted file mode 100644
index 594ab1d..0000000
--- a/tensorflow/compiler/tf2xla/lib/while_loop.cc
+++ /dev/null
@@ -1,127 +0,0 @@
-/* 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.
-==============================================================================*/
-
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
-#include "tensorflow/compiler/tf2xla/lib/util.h"
-#include "tensorflow/compiler/xla/client/xla_builder.h"
-#include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/status_macros.h"
-
-namespace tensorflow {
-
-xla::StatusOr<std::vector<xla::XlaOp>> XlaWhileLoop(
- const LoopConditionFunction& condition_function,
- const LoopBodyFunction& body_function,
- absl::Span<const xla::XlaOp> initial_values, absl::string_view name,
- xla::XlaBuilder* builder) {
- int arity = initial_values.size();
- std::vector<xla::Shape> var_shapes;
- var_shapes.reserve(arity);
- for (const xla::XlaOp& input : initial_values) {
- TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input));
- var_shapes.push_back(std::move(shape));
- }
- xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(var_shapes);
-
- // Unpacks a tuple into its component parts.
- auto unpack_tuple = [](xla::XlaOp tuple, int arity,
- xla::XlaBuilder* builder) {
- std::vector<xla::XlaOp> elements(arity);
- for (int i = 0; i < arity; ++i) {
- elements[i] = xla::GetTupleElement(tuple, i);
- }
- return elements;
- };
-
- // Build the condition.
- std::unique_ptr<xla::XlaBuilder> cond_builder =
- builder->CreateSubBuilder(absl::StrCat(name, "_condition"));
- {
- auto parameter =
- xla::Parameter(cond_builder.get(), 0, tuple_shape, "parameter");
-
- TF_RETURN_IF_ERROR(
- condition_function(unpack_tuple(parameter, arity, cond_builder.get()),
- cond_builder.get())
- .status());
- }
- TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build());
-
- // Build the body.
- std::unique_ptr<xla::XlaBuilder> body_builder =
- builder->CreateSubBuilder(absl::StrCat(name, "_body"));
- {
- auto parameter =
- xla::Parameter(body_builder.get(), 0, tuple_shape, "parameter");
-
- TF_ASSIGN_OR_RETURN(
- auto result,
- body_function(unpack_tuple(parameter, arity, body_builder.get()),
- body_builder.get()));
-
- TF_RET_CHECK(result.size() == initial_values.size());
- xla::Tuple(body_builder.get(), result);
- }
- TF_ASSIGN_OR_RETURN(auto body, body_builder->Build());
-
- auto outputs = xla::While(cond, body, xla::Tuple(builder, initial_values));
-
- return unpack_tuple(outputs, arity, builder);
-}
-
-xla::StatusOr<std::vector<xla::XlaOp>> XlaForEachIndex(
- int64 num_iterations, xla::PrimitiveType num_iterations_type,
- const ForEachIndexBodyFunction& body_function,
- absl::Span<const xla::XlaOp> initial_values, absl::string_view name,
- xla::XlaBuilder* builder) {
- auto while_cond_fn =
- [&](absl::Span<const xla::XlaOp> values,
- xla::XlaBuilder* cond_builder) -> xla::StatusOr<xla::XlaOp> {
- return xla::Lt(values[0], IntegerLiteral(cond_builder, num_iterations_type,
- num_iterations));
- };
- auto while_body_fn = [&](absl::Span<const xla::XlaOp> values,
- xla::XlaBuilder* body_builder)
- -> xla::StatusOr<std::vector<xla::XlaOp>> {
- xla::XlaOp iteration = values[0];
-
- std::vector<xla::XlaOp> updated_values;
- updated_values.reserve(values.size());
- updated_values.push_back(xla::Add(
- iteration,
- xla::ConstantLiteral(body_builder,
- xla::LiteralUtil::One(num_iterations_type))));
-
- values.remove_prefix(1);
- TF_ASSIGN_OR_RETURN(std::vector<xla::XlaOp> body_outputs,
- body_function(iteration, values, body_builder));
- updated_values.insert(updated_values.end(), body_outputs.begin(),
- body_outputs.end());
- return updated_values;
- };
-
- std::vector<xla::XlaOp> values;
- values.reserve(initial_values.size() + 1);
- values.push_back(xla::ConstantLiteral(
- builder, xla::LiteralUtil::Zero(num_iterations_type)));
- values.insert(values.end(), initial_values.begin(), initial_values.end());
-
- TF_ASSIGN_OR_RETURN(values, XlaWhileLoop(while_cond_fn, while_body_fn, values,
- name, builder));
- values.erase(values.begin(), values.begin() + 1);
- return values;
-}
-
-} // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/resource_operation_table.cc b/tensorflow/compiler/tf2xla/resource_operation_table.cc
index 72b2409..ff9f1b9 100644
--- a/tensorflow/compiler/tf2xla/resource_operation_table.cc
+++ b/tensorflow/compiler/tf2xla/resource_operation_table.cc
@@ -65,6 +65,7 @@
add("ResourceApplyFtrlV2" , kReadWrite, kVariable);
add("ResourceApplyGradientDescent" , kReadWrite, kVariable);
add("ResourceApplyMomentum" , kReadWrite, kVariable);
+ add("ResourceApplyKerasMomentum" , kReadWrite, kVariable);
add("ResourceApplyPowerSign" , kReadWrite, kVariable);
add("ResourceApplyProximalAdagrad" , kReadWrite, kVariable);
add("ResourceApplyProximalGradientDescent" , kReadWrite, kVariable);
diff --git a/tensorflow/compiler/tf2xla/side_effect_util.cc b/tensorflow/compiler/tf2xla/side_effect_util.cc
index b233e6b..b62f8e9 100644
--- a/tensorflow/compiler/tf2xla/side_effect_util.cc
+++ b/tensorflow/compiler/tf2xla/side_effect_util.cc
@@ -24,6 +24,8 @@
const char kXlaTokenArgNodeName[] = "_xla_token_arg_node";
+const char kXlaHasHostTransferAttrName[] = "_xla_has_host_transfer";
+
std::set<std::string> CalculateTokenInputsForOutputToken(const Graph& g) {
std::set<std::string> results;
Node* first_side_effecting_node_on_path = nullptr;
diff --git a/tensorflow/compiler/tf2xla/side_effect_util.h b/tensorflow/compiler/tf2xla/side_effect_util.h
index f22ddb2..7081b36 100644
--- a/tensorflow/compiler/tf2xla/side_effect_util.h
+++ b/tensorflow/compiler/tf2xla/side_effect_util.h
@@ -35,6 +35,9 @@
// node has side-effect dependency on current graph's token input.
extern const char kXlaTokenArgNodeName[];
+// This node have XlaRecvAtHost/XlaSendFromHost in its associated functions.
+extern const char kXlaHasHostTransferAttrName[];
+
// Calculates side-effect dependencies for the graph's token output.
// Returns a set of node names representing these dependencies.
std::set<std::string> CalculateTokenInputsForOutputToken(const Graph& g);
diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.cc b/tensorflow/compiler/tf2xla/tf2xla_util.cc
index cc81772..6cc8ae3 100644
--- a/tensorflow/compiler/tf2xla/tf2xla_util.cc
+++ b/tensorflow/compiler/tf2xla/tf2xla_util.cc
@@ -557,6 +557,12 @@
return true;
}
+ if (node_def.op() == "XlaHostCompute") {
+ // XlaHostCompute has "shape_inference_graph" func attr, but that's not
+ // related to graph execution.
+ return false;
+ }
+
for (const auto& iter : node_def.attr()) {
if (iter.second.has_func()) {
return true;
@@ -578,6 +584,9 @@
// This is a SymbolicGradient op.
AttrValueMap attrs(node.attrs().begin(), node.attrs().end());
results.emplace_back(AssociatedFunctionInfo::SymbolicGradient(op, attrs));
+ } else if (node.type_string() == "XlaHostCompute") {
+ // XlaHostCompute has "shape_inference_graph" func attr, but that's not
+ // related to graph execution.
} else {
// Collect all function attrs for the node.
for (auto& iter : node.attrs()) {
diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD
index 4360e08..0a20ddf 100644
--- a/tensorflow/compiler/xla/BUILD
+++ b/tensorflow/compiler/xla/BUILD
@@ -109,7 +109,7 @@
name = "status_macros",
srcs = ["status_macros.cc"],
hdrs = ["status_macros.h"],
- visibility = [":friends"],
+ visibility = ["//visibility:public"],
deps = [
":statusor",
":types",
@@ -224,6 +224,7 @@
name = "shape_util",
srcs = [
"index_util.cc",
+ "layout.cc",
"layout_util.cc",
"primitive_util.cc",
"shape.cc",
@@ -231,6 +232,7 @@
],
hdrs = [
"index_util.h",
+ "layout.h",
"layout_util.h",
"primitive_util.h",
"shape.h",
@@ -302,6 +304,22 @@
)
tf_cc_test(
+ name = "layout_test",
+ srcs = ["layout_test.cc"],
+ deps = [
+ ":shape_util",
+ ":status_macros",
+ ":test",
+ ":test_helpers",
+ ":types",
+ ":util",
+ ":xla_data_proto",
+ "//tensorflow/core:test_main",
+ "@com_google_absl//absl/strings",
+ ],
+)
+
+tf_cc_test(
name = "index_util_test",
srcs = ["index_util_test.cc"],
deps = [
@@ -705,7 +723,6 @@
"//tensorflow/compiler/xla/service:hlo",
"//tensorflow/compiler/xla/service:hlo_evaluator",
"//tensorflow/compiler/xla/service:shape_inference",
- "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_matmul",
"//tensorflow/core:lib",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/types:span",
diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD
index fe99564..e61d9d2 100644
--- a/tensorflow/compiler/xla/client/BUILD
+++ b/tensorflow/compiler/xla/client/BUILD
@@ -3,7 +3,7 @@
licenses(["notice"]) # Apache 2.0
-package(default_visibility = [":friends"])
+package(default_visibility = ["//visibility:public"])
package_group(
name = "friends",
diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc
index 74b76f9..43127ca 100644
--- a/tensorflow/compiler/xla/client/client.cc
+++ b/tensorflow/compiler/xla/client/client.cc
@@ -186,7 +186,7 @@
ComputeConstantGraphRequest request;
*request.mutable_computation() = computation.proto();
if (output_layout != nullptr) {
- *request.mutable_output_layout() = *output_layout;
+ *request.mutable_output_layout() = output_layout->ToProto();
}
ComputeConstantResponse response;
diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD
index f0f530d..970f007 100644
--- a/tensorflow/compiler/xla/client/lib/BUILD
+++ b/tensorflow/compiler/xla/client/lib/BUILD
@@ -1,5 +1,7 @@
# Common computation builders for XLA.
+load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites", "xla_test")
+
licenses(["notice"]) # Apache 2.0
package(default_visibility = ["//tensorflow/compiler/xla/client:friends"])
@@ -13,9 +15,6 @@
]),
)
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test")
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites")
-
# Generate test_suites for all backends, named "${backend}_tests".
generate_backend_suites()
@@ -36,6 +35,48 @@
)
cc_library(
+ name = "cholesky",
+ srcs = ["cholesky.cc"],
+ hdrs = ["cholesky.h"],
+ deps = [
+ ":math",
+ "//tensorflow/compiler/xla:literal",
+ "//tensorflow/compiler/xla:shape_util",
+ "//tensorflow/compiler/xla:status_macros",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/client/lib:constants",
+ "//tensorflow/compiler/xla/client/lib:loops",
+ "//tensorflow/compiler/xla/client/lib:matrix",
+ "//tensorflow/compiler/xla/client/lib:slicing",
+ "//tensorflow/compiler/xla/client/lib:triangular_solve",
+ "//tensorflow/core:lib",
+ ],
+)
+
+xla_test(
+ name = "cholesky_test",
+ srcs = ["cholesky_test.cc"],
+ tags = ["optonly"],
+ deps = [
+ ":arithmetic",
+ ":cholesky",
+ ":matrix",
+ "//tensorflow/compiler/xla:array2d",
+ "//tensorflow/compiler/xla:literal",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:test",
+ "//tensorflow/compiler/xla:types",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/tests:client_library_test_base",
+ "//tensorflow/compiler/xla/tests:literal_test_util",
+ "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+ "//tensorflow/core:test",
+ ],
+)
+
+cc_library(
name = "constants",
srcs = ["constants.cc"],
hdrs = ["constants.h"],
@@ -76,6 +117,22 @@
)
cc_library(
+ name = "loops",
+ srcs = ["loops.cc"],
+ hdrs = ["loops.h"],
+ deps = [
+ ":constants",
+ "//tensorflow/compiler/xla:shape_util",
+ "//tensorflow/compiler/xla:status_macros",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/client:xla_computation",
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/types:span",
+ ],
+)
+
+cc_library(
name = "math",
srcs = ["math.cc"],
hdrs = ["math.h"],
@@ -178,6 +235,48 @@
)
cc_library(
+ name = "qr",
+ srcs = ["qr.cc"],
+ hdrs = ["qr.h"],
+ deps = [
+ ":arithmetic",
+ ":constants",
+ ":loops",
+ ":math",
+ ":matrix",
+ ":slicing",
+ "//tensorflow/compiler/xla:literal_util",
+ "//tensorflow/compiler/xla:shape_util",
+ "//tensorflow/compiler/xla:status_macros",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/core:lib",
+ ],
+)
+
+xla_test(
+ name = "qr_test",
+ srcs = ["qr_test.cc"],
+ tags = ["optonly"],
+ deps = [
+ ":matrix",
+ ":qr",
+ "//tensorflow/compiler/xla:array2d",
+ "//tensorflow/compiler/xla:array3d",
+ "//tensorflow/compiler/xla:literal",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:test",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/tests:client_library_test_base",
+ "//tensorflow/compiler/xla/tests:literal_test_util",
+ "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+ "//tensorflow/core:test",
+ ],
+)
+
+cc_library(
name = "slicing",
srcs = ["slicing.cc"],
hdrs = ["slicing.h"],
@@ -238,6 +337,34 @@
)
cc_library(
+ name = "quantize",
+ hdrs = ["quantize.h"],
+ deps = [
+ ":constants",
+ "//tensorflow/compiler/xla:types",
+ "//tensorflow/compiler/xla:util",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/core:lib",
+ ],
+)
+
+xla_test(
+ name = "quantize_test",
+ srcs = ["quantize_test.cc"],
+ tags = ["enable_for_xla_interpreter"],
+ deps = [
+ ":quantize",
+ "//tensorflow/compiler/xla:test",
+ "//tensorflow/compiler/xla:types",
+ "//tensorflow/compiler/xla:util",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/tests:client_library_test_base",
+ "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+ ],
+)
+
+cc_library(
name = "testing",
srcs = ["testing.cc"],
hdrs = ["testing.h"],
@@ -258,3 +385,48 @@
"@com_google_absl//absl/strings",
],
)
+
+cc_library(
+ name = "triangular_solve",
+ srcs = ["triangular_solve.cc"],
+ hdrs = ["triangular_solve.h"],
+ deps = [
+ "//tensorflow/compiler/xla:literal",
+ "//tensorflow/compiler/xla:shape_util",
+ "//tensorflow/compiler/xla:status_macros",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:util",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/client:xla_computation",
+ "//tensorflow/compiler/xla/client/lib:constants",
+ "//tensorflow/compiler/xla/client/lib:math",
+ "//tensorflow/compiler/xla/client/lib:matrix",
+ "//tensorflow/compiler/xla/client/lib:slicing",
+ "//tensorflow/core:lib",
+ ],
+)
+
+xla_test(
+ name = "triangular_solve_test",
+ srcs = ["triangular_solve_test.cc"],
+ tags = ["noasan"], # sometimes times out, http://b/78650012
+ deps = [
+ ":triangular_solve",
+ "//tensorflow/compiler/xla:array2d",
+ "//tensorflow/compiler/xla:literal",
+ "//tensorflow/compiler/xla:shape_util",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:test",
+ "//tensorflow/compiler/xla:types",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:global_data",
+ "//tensorflow/compiler/xla/client:local_client",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/tests:client_library_test_base",
+ "//tensorflow/compiler/xla/tests:literal_test_util",
+ "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+ "//tensorflow/core:lib",
+ "//tensorflow/core:test",
+ ],
+)
diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.cc b/tensorflow/compiler/xla/client/lib/cholesky.cc
similarity index 67%
rename from tensorflow/compiler/tf2xla/lib/cholesky.cc
rename to tensorflow/compiler/xla/client/lib/cholesky.cc
index 7ef8659..fd98049 100644
--- a/tensorflow/compiler/tf2xla/lib/cholesky.cc
+++ b/tensorflow/compiler/xla/client/lib/cholesky.cc
@@ -13,17 +13,17 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/cholesky.h"
+#include "tensorflow/compiler/xla/client/lib/cholesky.h"
#include <memory>
#include <vector>
-#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h"
-#include "tensorflow/compiler/tf2xla/lib/util.h"
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
+#include "tensorflow/compiler/xla/client/lib/loops.h"
+#include "tensorflow/compiler/xla/client/lib/math.h"
#include "tensorflow/compiler/xla/client/lib/matrix.h"
#include "tensorflow/compiler/xla/client/lib/slicing.h"
+#include "tensorflow/compiler/xla/client/lib/triangular_solve.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/shape_util.h"
@@ -31,7 +31,7 @@
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/core/lib/core/errors.h"
-namespace tensorflow {
+namespace xla {
namespace {
@@ -50,26 +50,25 @@
// l[..., j+1:, j] = (a[..., j+1:, j] - np.dot(l[..., j+1:, :j], row_t)) /
// l[..., j, j]
// return l
-xla::XlaOp CholeskyUnblocked(xla::XlaOp a,
- xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = a.builder();
- return builder->ReportErrorOrReturn([&]() -> xla::StatusOr<xla::XlaOp> {
- TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
- const int n_dims = xla::ShapeUtil::Rank(a_shape);
- const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
- auto major_dims = xla::AsInt64Slice(a_shape.dimensions())
+XlaOp CholeskyUnblocked(XlaOp a, PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = a.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
+ const int n_dims = ShapeUtil::Rank(a_shape);
+ const int64 n = ShapeUtil::GetDimension(a_shape, -1);
+ auto major_dims = AsInt64Slice(a_shape.dimensions())
.subspan(
/*pos=*/0,
/*len=*/n_dims - 2);
- xla::XlaOp l = xla::ZerosLike(a);
+ XlaOp l = ZerosLike(a);
// Construct the for loop body to iterate over rows.
- auto body_fn = [&](xla::XlaOp i, absl::Span<const xla::XlaOp> loop_vars,
- xla::XlaBuilder* body_builder)
- -> xla::StatusOr<std::vector<xla::XlaOp>> {
- xla::Shape col_shape;
- xla::Shape row_shape;
+ auto body_fn =
+ [&](XlaOp i, absl::Span<const XlaOp> loop_vars,
+ XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> {
+ Shape col_shape;
+ Shape row_shape;
for (int64 d : major_dims) {
row_shape.add_dimensions(d);
col_shape.add_dimensions(d);
@@ -77,43 +76,40 @@
row_shape.add_dimensions(1);
row_shape.add_dimensions(n);
row_shape.set_element_type(a_shape.element_type());
- auto mask_zeros_row = xla::Zeros(body_builder, row_shape);
+ auto mask_zeros_row = Zeros(body_builder, row_shape);
col_shape.add_dimensions(n);
col_shape.add_dimensions(1);
col_shape.set_element_type(a_shape.element_type());
- auto mask_zeros_col = xla::Zeros(body_builder, col_shape);
+ auto mask_zeros_col = Zeros(body_builder, col_shape);
std::vector<int32> mask_vector(n);
std::iota(mask_vector.begin(), mask_vector.end(), 0);
- auto mask_range = xla::ConstantR1<int32>(body_builder, mask_vector);
+ auto mask_range = ConstantR1<int32>(body_builder, mask_vector);
auto mask_range_row =
- xla::Broadcast(xla::Reshape(mask_range, {0}, {1, n}), major_dims);
+ Broadcast(Reshape(mask_range, {0}, {1, n}), major_dims);
auto mask_range_col =
- xla::Broadcast(xla::Reshape(mask_range, {0}, {n, 1}), major_dims);
+ Broadcast(Reshape(mask_range, {0}, {n, 1}), major_dims);
auto body_a = loop_vars[0];
auto body_l = loop_vars[1];
// row = l[..., i, :i]
// select the whole i-th row, then mask out all columns past i-1
- auto zero = xla::ConstantR0<int32>(body_builder, 0);
+ auto zero = ConstantR0<int32>(body_builder, 0);
auto l_i = DynamicSliceInMinorDims(body_l, {i, zero}, {1, n});
- auto row = xla::Select(xla::Ge(mask_range_row, i), mask_zeros_row, l_i);
+ auto row = Select(Ge(mask_range_row, i), mask_zeros_row, l_i);
// a[..., i, i]
auto a_ii = DynamicSliceInMinorDims(body_a, {i, i}, {1, 1});
// np.dot(row, np.swapaxes(row, -1, -2))
auto diag_dot = BatchDot(row, TransposeInMinorDims(row), precision);
// l[..., i, i] = np.sqrt(a[..., i, i] - np.dot(row,
// np.swapaxes(row, -1, -2)))
- auto l_ii =
- xla::Pow(a_ii - diag_dot,
- FloatLiteral(body_builder, a_shape.element_type(), 0.5));
+ auto l_ii = Sqrt(a_ii - diag_dot);
// a[..., i+1:, i]
// select the whole i-th column, then mask out all rows above i+1
auto a_0i = DynamicSliceInMinorDims(body_a, {i}, {1});
- auto a_ip1i =
- xla::Select(xla::Le(mask_range_col, i), mask_zeros_col, a_0i);
+ auto a_ip1i = Select(Le(mask_range_col, i), mask_zeros_col, a_0i);
// l[..., i+1:, i] = (a[..., i+1:, i] - np.dot(l[..., i+1:, :i], r.T)) /
// l[..., i, i]
@@ -122,8 +118,7 @@
// r.T)
auto dot = BatchDot(body_l, TransposeInMinorDims(row), precision);
// np.dot(l[..., i+1:, :i], r.T)
- auto dot_ip1 =
- xla::Select(xla::Le(mask_range_col, i), mask_zeros_col, dot);
+ auto dot_ip1 = Select(Le(mask_range_col, i), mask_zeros_col, dot);
body_l =
DynamicUpdateSliceInMinorDims(body_l, (a_ip1i - dot_ip1) / l_ii, {i});
@@ -131,12 +126,12 @@
// column assign will wrap around and overwrite the diagonal assign.
body_l = DynamicUpdateSliceInMinorDims(body_l, l_ii, {i, i});
- return std::vector<xla::XlaOp>{body_a, body_l};
+ return std::vector<XlaOp>{body_a, body_l};
};
TF_ASSIGN_OR_RETURN(
auto cholesky_while,
- XlaForEachIndex(n, xla::S32, body_fn, {a, l}, "unblocked", builder));
+ ForEachIndex(n, S32, body_fn, {a, l}, "unblocked", builder));
return cholesky_while[1];
});
@@ -144,34 +139,35 @@
} // namespace
-xla::XlaOp Cholesky(xla::XlaOp a, int64 block_size,
- xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = a.builder();
- return builder->ReportErrorOrReturn([&]() -> xla::StatusOr<xla::XlaOp> {
- TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
- const int ndims = xla::ShapeUtil::Rank(a_shape);
+XlaOp Cholesky(XlaOp a, int64 block_size,
+ PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = a.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
+ const int ndims = ShapeUtil::Rank(a_shape);
if (ndims < 2) {
- return errors::InvalidArgument(
- "Arguments to Cholesky must have rank >= 2: ", ndims);
+ return InvalidArgument(
+ "Argument to Cholesky must have rank >= 2; shape was %s",
+ a_shape.ToString());
}
- const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
- if (n != xla::ShapeUtil::GetDimension(a_shape, -2)) {
- return errors::InvalidArgument(
- "Arguments to Cholesky must be square matrices: ",
- xla::ShapeUtil::HumanString(a_shape));
+ const int64 n = ShapeUtil::GetDimension(a_shape, -1);
+ if (n != ShapeUtil::GetDimension(a_shape, -2)) {
+ return InvalidArgument(
+ "Argument to Cholesky must be batched square matrices; got shape %s",
+ ShapeUtil::HumanString(a_shape));
}
if (block_size < 1) {
- return errors::InvalidArgument(
- "block_size argument to Cholesky must be >= 1; got ", block_size);
+ return InvalidArgument(
+ "block_size argument to Cholesky must be >= 1; got %d", block_size);
}
// Blocked left-looking Cholesky factorization.
// Algorithm 1 from
// Haidar, Azzam, et al. "High-performance Cholesky factorization for
// GPU-only execution." Proceedings of General Purpose GPUs. ACM, 2017.
- xla::XlaOp l = xla::ZerosLike(a);
+ XlaOp l = ZerosLike(a);
for (int64 i = 0; i < n; i += block_size) {
int64 k = std::min(block_size, n - i);
if (i > 0) {
@@ -207,4 +203,4 @@
});
}
-} // namespace tensorflow
+} // namespace xla
diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.h b/tensorflow/compiler/xla/client/lib/cholesky.h
similarity index 86%
rename from tensorflow/compiler/tf2xla/lib/cholesky.h
rename to tensorflow/compiler/xla/client/lib/cholesky.h
index 9a561c3..0bae268 100644
--- a/tensorflow/compiler/tf2xla/lib/cholesky.h
+++ b/tensorflow/compiler/xla/client/lib/cholesky.h
@@ -13,13 +13,13 @@
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_
-#define TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_
+#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_
+#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
-namespace tensorflow {
+namespace xla {
// Computes the Cholesky decompositions of a batch of symmetric positive
// definite matrices.
@@ -34,6 +34,6 @@
xla::XlaOp a, int64 block_size = 256,
xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST);
-} // namespace tensorflow
+} // namespace xla
-#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_
+#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_
diff --git a/tensorflow/compiler/xla/client/lib/cholesky_test.cc b/tensorflow/compiler/xla/client/lib/cholesky_test.cc
new file mode 100644
index 0000000..ba9580a
--- /dev/null
+++ b/tensorflow/compiler/xla/client/lib/cholesky_test.cc
@@ -0,0 +1,166 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/client/lib/cholesky.h"
+
+#include <memory>
+#include <numeric>
+#include <vector>
+
+#include "tensorflow/compiler/xla/array2d.h"
+#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
+#include "tensorflow/compiler/xla/client/lib/matrix.h"
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/compiler/xla/literal.h"
+#include "tensorflow/compiler/xla/statusor.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
+#include "tensorflow/compiler/xla/tests/literal_test_util.h"
+#include "tensorflow/compiler/xla/tests/test_macros.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace {
+
+using xla::int64;
+
+using CholeskyTest = xla::ClientLibraryTestBase;
+
+XLA_TEST_F(CholeskyTest, Simple) {
+ xla::XlaBuilder builder(TestName());
+
+ xla::Array2D<float> a_vals({
+ {4, 6, 8, 10},
+ {6, 45, 54, 63},
+ {8, 54, 146, 166},
+ {10, 63, 166, 310},
+ });
+
+ xla::XlaOp a;
+ auto a_data = CreateR2Parameter<float>(a_vals, 0, "a", &builder, &a);
+ xla::Cholesky(a, /*block_size=*/2);
+
+ xla::Array2D<float> expected({
+ {2, 0, 0, 0},
+ {3, 6, 0, 0},
+ {4, 7, 9, 0},
+ {5, 8, 10, 11},
+ });
+
+ ComputeAndCompareR2<float>(&builder, expected, {a_data.get()},
+ xla::ErrorSpec(1e-4, 1e-4));
+}
+
+XLA_TEST_F(CholeskyTest, Simple2) {
+ xla::XlaBuilder builder(TestName());
+
+ xla::Array2D<float> a_vals({
+ {16, 24, 8, 12},
+ {24, 61, 82, 48},
+ {8, 82, 456, 106},
+ {12, 48, 106, 62},
+ });
+
+ xla::XlaOp a;
+ auto a_data = CreateR2Parameter<float>(a_vals, 0, "a", &builder, &a);
+ xla::Cholesky(a);
+
+ xla::Array2D<float> expected(
+ {{4, 0, 0, 0}, {6, 5, 0, 0}, {2, 14, 16, 0}, {3, 6, 1, 4}});
+
+ ComputeAndCompareR2<float>(&builder, expected, {a_data.get()},
+ xla::ErrorSpec(1e-4, 1e-4));
+}
+
+XLA_TEST_F(CholeskyTest, SimpleBatched) {
+ xla::XlaBuilder builder(TestName());
+
+ xla::Array3D<float> a_vals({
+ {
+ {4, 6, 8, 10},
+ {6, 45, 54, 63},
+ {8, 54, 146, 166},
+ {10, 63, 166, 310},
+ },
+ {
+ {16, 24, 8, 12},
+ {24, 61, 82, 48},
+ {8, 82, 456, 106},
+ {12, 48, 106, 62},
+ },
+ });
+
+ xla::XlaOp a;
+ auto a_data = CreateR3Parameter<float>(a_vals, 0, "a", &builder, &a);
+ xla::Cholesky(a);
+
+ xla::Array3D<float> expected({
+ {
+ {2, 0, 0, 0},
+ {3, 6, 0, 0},
+ {4, 7, 9, 0},
+ {5, 8, 10, 11},
+ },
+ {{4, 0, 0, 0}, {6, 5, 0, 0}, {2, 14, 16, 0}, {3, 6, 1, 4}},
+ });
+
+ ComputeAndCompareR3<float>(&builder, expected, {a_data.get()},
+ xla::ErrorSpec(1e-4, 1e-4));
+}
+
+using CholeskyTestCase = std::tuple<int64, int64>;
+
+class RandomCholeskyTest
+ : public xla::ClientLibraryTestBase,
+ public ::testing::WithParamInterface<CholeskyTestCase> {};
+
+XLA_TEST_P(RandomCholeskyTest, Random) {
+ xla::XlaBuilder builder(TestName());
+
+ auto test_params = GetParam();
+ std::vector<int64> dimensions = {std::get<0>(test_params),
+ std::get<1>(test_params),
+ std::get<1>(test_params)};
+ xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, dimensions);
+ TF_ASSERT_OK_AND_ASSIGN(
+ auto literal,
+ xla::LiteralUtil::CreateRandomLiteral<xla::F32>(shape, 0.0, 1.0));
+
+ auto input = xla::Parameter(&builder, 0, shape, "input");
+ // Form a random positive definite matrix.
+ auto matrix = xla::BatchDot(input, TransposeInMinorDims(input),
+ xla::PrecisionConfig::HIGHEST);
+
+ auto cholesky = xla::Cholesky(matrix, /*block_size=*/4);
+
+ // Verify that ||matrix - cholesky * cholesky_t||_2 ~= 0
+ auto verification = xla::BatchDot(cholesky, TransposeInMinorDims(cholesky),
+ xla::PrecisionConfig::HIGHEST);
+ auto delta = matrix - verification;
+ xla::Reduce(delta * delta, xla::ConstantR0<float>(&builder, 0.0),
+ CreateScalarAddComputation(xla::F32, &builder), {0, 1, 2});
+
+ TF_ASSERT_OK_AND_ASSIGN(auto input_data, client_->TransferToServer(literal));
+ ComputeAndCompareR0<float>(&builder, 0.0, {input_data.get()},
+ xla::ErrorSpec(1e-4, 1e-4));
+}
+
+INSTANTIATE_TEST_CASE_P(RandomCholeskyTestInstance, RandomCholeskyTest,
+ ::testing::Values(CholeskyTestCase{1, 1},
+ CholeskyTestCase{1, 2},
+ CholeskyTestCase{10, 5},
+ CholeskyTestCase{2, 20}));
+
+} // namespace
diff --git a/tensorflow/compiler/xla/client/lib/loops.cc b/tensorflow/compiler/xla/client/lib/loops.cc
new file mode 100644
index 0000000..721f987
--- /dev/null
+++ b/tensorflow/compiler/xla/client/lib/loops.cc
@@ -0,0 +1,123 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/client/lib/loops.h"
+
+#include "tensorflow/compiler/xla/client/lib/constants.h"
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/status_macros.h"
+
+namespace xla {
+
+StatusOr<std::vector<XlaOp>> WhileLoopHelper(
+ const WhileLoopHelperConditionFunction& condition_function,
+ const WhileLoopHelperBodyFunction& body_function,
+ absl::Span<const XlaOp> initial_values, absl::string_view name,
+ XlaBuilder* builder) {
+ int arity = initial_values.size();
+ std::vector<Shape> var_shapes;
+ var_shapes.reserve(arity);
+ for (const XlaOp& input : initial_values) {
+ TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input));
+ var_shapes.push_back(std::move(shape));
+ }
+ Shape tuple_shape = ShapeUtil::MakeTupleShape(var_shapes);
+
+ // Unpacks a tuple into its component parts.
+ auto unpack_tuple = [](XlaOp tuple, int arity, XlaBuilder* builder) {
+ std::vector<XlaOp> elements(arity);
+ for (int i = 0; i < arity; ++i) {
+ elements[i] = GetTupleElement(tuple, i);
+ }
+ return elements;
+ };
+
+ // Build the condition.
+ std::unique_ptr<XlaBuilder> cond_builder =
+ builder->CreateSubBuilder(absl::StrCat(name, "_condition"));
+ {
+ auto parameter = Parameter(cond_builder.get(), 0, tuple_shape, "parameter");
+
+ TF_RETURN_IF_ERROR(
+ condition_function(unpack_tuple(parameter, arity, cond_builder.get()),
+ cond_builder.get())
+ .status());
+ }
+ TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build());
+
+ // Build the body.
+ std::unique_ptr<XlaBuilder> body_builder =
+ builder->CreateSubBuilder(absl::StrCat(name, "_body"));
+ {
+ auto parameter = Parameter(body_builder.get(), 0, tuple_shape, "parameter");
+
+ TF_ASSIGN_OR_RETURN(
+ auto result,
+ body_function(unpack_tuple(parameter, arity, body_builder.get()),
+ body_builder.get()));
+
+ TF_RET_CHECK(result.size() == initial_values.size());
+ Tuple(body_builder.get(), result);
+ }
+ TF_ASSIGN_OR_RETURN(auto body, body_builder->Build());
+
+ auto outputs = While(cond, body, Tuple(builder, initial_values));
+
+ return unpack_tuple(outputs, arity, builder);
+}
+
+StatusOr<std::vector<XlaOp>> ForEachIndex(
+ int64 num_iterations, PrimitiveType num_iterations_type,
+ const ForEachIndexBodyFunction& body_function,
+ absl::Span<const XlaOp> initial_values, absl::string_view name,
+ XlaBuilder* builder) {
+ auto while_cond_fn = [&](absl::Span<const XlaOp> values,
+ XlaBuilder* cond_builder) -> StatusOr<XlaOp> {
+ return Lt(values[0], ConstantR0WithType(cond_builder, num_iterations_type,
+ num_iterations));
+ };
+ auto while_body_fn =
+ [&](absl::Span<const XlaOp> values,
+ XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> {
+ XlaOp iteration = values[0];
+
+ std::vector<XlaOp> updated_values;
+ updated_values.reserve(values.size());
+ updated_values.push_back(Add(
+ iteration,
+ ConstantLiteral(body_builder, LiteralUtil::One(num_iterations_type))));
+
+ values.remove_prefix(1);
+ TF_ASSIGN_OR_RETURN(std::vector<XlaOp> body_outputs,
+ body_function(iteration, values, body_builder));
+ updated_values.insert(updated_values.end(), body_outputs.begin(),
+ body_outputs.end());
+ return updated_values;
+ };
+
+ std::vector<XlaOp> values;
+ values.reserve(initial_values.size() + 1);
+ values.push_back(
+ ConstantLiteral(builder, LiteralUtil::Zero(num_iterations_type)));
+ values.insert(values.end(), initial_values.begin(), initial_values.end());
+
+ TF_ASSIGN_OR_RETURN(values, WhileLoopHelper(while_cond_fn, while_body_fn,
+ values, name, builder));
+ values.erase(values.begin(), values.begin() + 1);
+ return values;
+}
+
+} // namespace xla
diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.h b/tensorflow/compiler/xla/client/lib/loops.h
similarity index 61%
rename from tensorflow/compiler/tf2xla/lib/while_loop.h
rename to tensorflow/compiler/xla/client/lib/loops.h
index f2134bb..e11de59 100644
--- a/tensorflow/compiler/tf2xla/lib/while_loop.h
+++ b/tensorflow/compiler/xla/client/lib/loops.h
@@ -13,8 +13,8 @@
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_
-#define TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_
+#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_LOOPS_H_
+#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_LOOPS_H_
#include <functional>
#include <vector>
@@ -25,19 +25,18 @@
#include "tensorflow/compiler/xla/client/xla_computation.h"
#include "tensorflow/compiler/xla/statusor.h"
-namespace tensorflow {
+namespace xla {
// Function that builds a loop condition. Takes as input a sequence of input
// values, and returns a boolean value representing if the condition succeeds.
-typedef std::function<xla::StatusOr<xla::XlaOp>(absl::Span<const xla::XlaOp>,
- xla::XlaBuilder*)>
- LoopConditionFunction;
+typedef std::function<StatusOr<XlaOp>(absl::Span<const XlaOp>, XlaBuilder*)>
+ WhileLoopHelperConditionFunction;
// Function that builds a loop body. Takes as input a sequence of input values
// and returns a sequence of output values.
-typedef std::function<xla::StatusOr<std::vector<xla::XlaOp>>(
- absl::Span<const xla::XlaOp>, xla::XlaBuilder*)>
- LoopBodyFunction;
+typedef std::function<StatusOr<std::vector<XlaOp>>(absl::Span<const XlaOp>,
+ XlaBuilder*)>
+ WhileLoopHelperBodyFunction;
// Helper function for building an XLA while loop, where the values carried by
// the loop are a tuple of values, e.g., (a, b, c):
@@ -47,27 +46,27 @@
// init: (a, b, c)
// )
// 'name' is a descriptive name for the loop.
-xla::StatusOr<std::vector<xla::XlaOp>> XlaWhileLoop(
- const LoopConditionFunction& condition_function,
- const LoopBodyFunction& body_function,
- absl::Span<const xla::XlaOp> initial_values, absl::string_view name,
- xla::XlaBuilder* builder);
+StatusOr<std::vector<XlaOp>> WhileLoopHelper(
+ const WhileLoopHelperConditionFunction& condition_function,
+ const WhileLoopHelperBodyFunction& body_function,
+ absl::Span<const XlaOp> initial_values, absl::string_view name,
+ XlaBuilder* builder);
// Builds an XLA loop that repeats a computation `num_iterations` times.
//
// The body function (ForEachIndexBodyFunction) takes as input a pair of
// (current iteration number, loop-carried values), and returns an updated
// vector of the loop-carried values.
-typedef std::function<xla::StatusOr<std::vector<xla::XlaOp>>(
- xla::XlaOp, absl::Span<const xla::XlaOp>, xla::XlaBuilder*)>
+typedef std::function<StatusOr<std::vector<XlaOp>>(
+ XlaOp, absl::Span<const XlaOp>, XlaBuilder*)>
ForEachIndexBodyFunction;
-xla::StatusOr<std::vector<xla::XlaOp>> XlaForEachIndex(
- int64 num_iterations, xla::PrimitiveType num_iterations_type,
+StatusOr<std::vector<XlaOp>> ForEachIndex(
+ int64 num_iterations, PrimitiveType num_iterations_type,
const ForEachIndexBodyFunction& body_function,
- absl::Span<const xla::XlaOp> initial_values, absl::string_view name,
- xla::XlaBuilder* builder);
+ absl::Span<const XlaOp> initial_values, absl::string_view name,
+ XlaBuilder* builder);
-} // namespace tensorflow
+} // namespace xla
-#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_
+#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_LOOPS_H_
diff --git a/tensorflow/compiler/tf2xla/lib/qr.cc b/tensorflow/compiler/xla/client/lib/qr.cc
similarity index 62%
rename from tensorflow/compiler/tf2xla/lib/qr.cc
rename to tensorflow/compiler/xla/client/lib/qr.cc
index d600774..72ca653 100644
--- a/tensorflow/compiler/tf2xla/lib/qr.cc
+++ b/tensorflow/compiler/xla/client/lib/qr.cc
@@ -13,15 +13,14 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/qr.h"
+#include "tensorflow/compiler/xla/client/lib/qr.h"
#include <memory>
#include <vector>
-#include "tensorflow/compiler/tf2xla/lib/util.h"
-#include "tensorflow/compiler/tf2xla/lib/while_loop.h"
#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
+#include "tensorflow/compiler/xla/client/lib/loops.h"
#include "tensorflow/compiler/xla/client/lib/math.h"
#include "tensorflow/compiler/xla/client/lib/matrix.h"
#include "tensorflow/compiler/xla/client/lib/slicing.h"
@@ -32,10 +31,18 @@
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/core/lib/core/errors.h"
-namespace tensorflow {
+namespace xla {
namespace {
+std::vector<int64> ConcatVectors(absl::Span<const int64> xs,
+ absl::Span<const int64> ys) {
+ std::vector<int64> output(xs.size() + ys.size());
+ std::copy(xs.begin(), xs.end(), output.begin());
+ std::copy(ys.begin(), ys.end(), output.begin() + xs.size());
+ return output;
+}
+
// Computes a Householder reflection of the form:
// H = I - tau v v.T.
// such that
@@ -65,52 +72,47 @@
// return (v, tau, beta)
// TODO(phawkins): LAPACK's xLARFG implementation has code for handling
// overflows in the norm/beta calculations. Perhaps do the same here.
-xla::Status House(xla::XlaOp x, xla::XlaOp k,
- absl::Span<const int64> batch_dims, const int64 m,
- xla::XlaOp* v, xla::XlaOp* tau, xla::XlaOp* beta) {
- xla::XlaBuilder* const builder = x.builder();
- TF_ASSIGN_OR_RETURN(xla::Shape x_shape, builder->GetShape(x));
- const xla::PrimitiveType type = x_shape.element_type();
+Status House(XlaOp x, XlaOp k, absl::Span<const int64> batch_dims,
+ const int64 m, XlaOp* v, XlaOp* tau, XlaOp* beta) {
+ XlaBuilder* const builder = x.builder();
+ TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x));
+ const PrimitiveType type = x_shape.element_type();
std::vector<int64> batch_dim_ids(batch_dims.size());
std::iota(batch_dim_ids.begin(), batch_dim_ids.end(), 0);
const int64 minor_dim = batch_dims.size();
- xla::XlaOp zero = xla::ScalarLike(x, 0.0);
- xla::XlaOp one = xla::ScalarLike(x, 1.0);
+ XlaOp zero = ScalarLike(x, 0.0);
+ XlaOp one = ScalarLike(x, 1.0);
// alpha = x[k]
- xla::XlaOp alpha =
- xla::Reshape(DynamicSliceInMinorDims(x, {k}, {1}), batch_dims);
+ XlaOp alpha = Reshape(DynamicSliceInMinorDims(x, {k}, {1}), batch_dims);
// Compute x[k+1:] (padded with zeros in elements 0..k)
- xla::XlaOp iota = xla::Iota(builder, xla::S32, m);
- xla::XlaOp x_after_k =
- xla::Mul(x, xla::ConvertElementType(xla::Gt(iota, k), type),
- /*broadcast_dimensions=*/{minor_dim});
+ XlaOp iota = Iota(builder, S32, m);
+ XlaOp x_after_k = Mul(x, ConvertElementType(Gt(iota, k), type),
+ /*broadcast_dimensions=*/{minor_dim});
// sigma = np.dot(x[k+1:], x[k+1:])
- auto sigma =
- xla::Reduce(x_after_k * x_after_k, zero,
- xla::CreateScalarAddComputation(type, builder), {minor_dim});
+ auto sigma = Reduce(x_after_k * x_after_k, zero,
+ CreateScalarAddComputation(type, builder), {minor_dim});
// mu = np.sqrt(x[k]*x[k] + sigma)
- auto mu = xla::Sqrt(xla::Square(alpha) + sigma);
+ auto mu = Sqrt(Square(alpha) + sigma);
- auto sigma_is_zero = xla::Eq(sigma, zero);
+ auto sigma_is_zero = Eq(sigma, zero);
- *beta = xla::Select(sigma_is_zero, alpha, -xla::Sign(alpha) * mu);
- *tau = xla::Select(sigma_is_zero, xla::Broadcast(zero, batch_dims),
- (*beta - alpha) / *beta);
- auto divisor = xla::Select(sigma_is_zero, xla::Broadcast(one, batch_dims),
- alpha - *beta);
+ *beta = Select(sigma_is_zero, alpha, -Sign(alpha) * mu);
+ *tau = Select(sigma_is_zero, Broadcast(zero, batch_dims),
+ (*beta - alpha) / *beta);
+ auto divisor =
+ Select(sigma_is_zero, Broadcast(one, batch_dims), alpha - *beta);
- auto e_k = xla::Broadcast(xla::ConvertElementType(xla::Eq(iota, k), type),
- std::vector<int64>(batch_dims.size(), 1));
+ auto e_k = Broadcast(ConvertElementType(Eq(iota, k), type),
+ std::vector<int64>(batch_dims.size(), 1));
// Form v as [0, 0, ..., 1] ++ x[k+1:] / divisor
// If sigma is zero, x[k+1:] is zero, so use any non-zero divisor.
- *v = e_k +
- xla::Div(x_after_k, divisor, /*broadcast_dimensions=*/batch_dim_ids);
+ *v = e_k + Div(x_after_k, divisor, /*broadcast_dimensions=*/batch_dim_ids);
return Status::OK();
}
@@ -143,90 +145,86 @@
// return (q, vs, taus)
struct QRBlockResult {
// The factored R value
- xla::XlaOp r;
+ XlaOp r;
// Representation of the Householder matrices I - beta v v.T
- xla::XlaOp taus; // Shape: [..., n]
- xla::XlaOp vs; // Shape: [..., m, n]
+ XlaOp taus; // Shape: [..., n]
+ XlaOp vs; // Shape: [..., m, n]
};
-xla::StatusOr<QRBlockResult> QRBlock(
- xla::XlaOp a, xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = a.builder();
- TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
- const int num_dims = xla::ShapeUtil::Rank(a_shape);
+StatusOr<QRBlockResult> QRBlock(XlaOp a, PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = a.builder();
+ TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
+ const int num_dims = ShapeUtil::Rank(a_shape);
if (num_dims < 2) {
- return errors::InvalidArgument("Arguments to QR must have rank >= 2: ",
- num_dims);
+ return InvalidArgument("Argument to QR must have rank >= 2; got shape %s",
+ a_shape.ToString());
}
- xla::PrimitiveType type = a_shape.element_type();
+ PrimitiveType type = a_shape.element_type();
- const int64 m = xla::ShapeUtil::GetDimension(a_shape, -2);
- const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
+ const int64 m = ShapeUtil::GetDimension(a_shape, -2);
+ const int64 n = ShapeUtil::GetDimension(a_shape, -1);
const int64 num_batch_dims = num_dims - 2;
std::vector<int64> batch_dims(num_batch_dims);
for (int i = 0; i < num_batch_dims; ++i) {
- batch_dims[i] = xla::ShapeUtil::GetDimension(a_shape, i);
+ batch_dims[i] = ShapeUtil::GetDimension(a_shape, i);
}
std::vector<int64> batch_dim_indices(num_batch_dims);
std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0);
- auto qr_body_fn =
- [&](xla::XlaOp j, absl::Span<const xla::XlaOp> values,
- xla::XlaBuilder* builder) -> xla::StatusOr<std::vector<xla::XlaOp>> {
+ auto qr_body_fn = [&](XlaOp j, absl::Span<const XlaOp> values,
+ XlaBuilder* builder) -> StatusOr<std::vector<XlaOp>> {
auto a = values[0];
auto vs = values[1];
auto taus = values[2];
// v, beta = house(a[:, j], j)
auto x = DynamicSliceInMinorDims(a, {j}, {1});
- xla::XlaOp v, tau, beta;
- TF_RETURN_IF_ERROR(House(xla::Collapse(x, {num_dims - 2, num_dims - 1}), j,
+ XlaOp v, tau, beta;
+ TF_RETURN_IF_ERROR(House(Collapse(x, {num_dims - 2, num_dims - 1}), j,
batch_dims, m, &v, &tau, &beta));
std::vector<int64> shape = batch_dims;
shape.push_back(1);
shape.push_back(m);
- auto v_broadcast = xla::Reshape(v, shape);
+ auto v_broadcast = Reshape(v, shape);
// a[:, :] -= tau * np.dot(v[:, np.newaxis],
// np.dot(v[np.newaxis, :], a[:, :]))
auto vva = BatchDot(v_broadcast, a, precision);
vva = BatchDot(TransposeInMinorDims(v_broadcast), vva, precision);
- a = a - xla::Mul(tau, vva,
- /*broadcast_dimensions=*/batch_dim_indices);
+ a = a - Mul(tau, vva,
+ /*broadcast_dimensions=*/batch_dim_indices);
// It is more precise to populate column 'k' explicitly, rather than
// computing it implicitly by applying the Householder transformation.
// a[k,k] = beta
// a[k+1:,k] = np.zeros([m-k-1], dtype=a.dtype)
- auto iota = xla::Reshape(xla::Iota(a.builder(), xla::S32, m), {m, 1});
- auto predecessor_mask = xla::ConvertElementType(xla::Lt(iota, j), type);
- auto mask = xla::Broadcast(xla::ConvertElementType(xla::Eq(iota, j), type),
- std::vector<int64>(batch_dims.size(), 1));
- auto new_x =
- xla::Mul(x, predecessor_mask,
- /*broadcast_dimensions=*/{num_dims - 2, num_dims - 1}) +
- xla::Mul(beta, mask, /*broadcast_dimensions=*/batch_dim_indices);
+ auto iota = Reshape(Iota(a.builder(), S32, m), {m, 1});
+ auto predecessor_mask = ConvertElementType(Lt(iota, j), type);
+ auto mask = Broadcast(ConvertElementType(Eq(iota, j), type),
+ std::vector<int64>(batch_dims.size(), 1));
+ auto new_x = Mul(x, predecessor_mask,
+ /*broadcast_dimensions=*/{num_dims - 2, num_dims - 1}) +
+ Mul(beta, mask, /*broadcast_dimensions=*/batch_dim_indices);
a = DynamicUpdateSliceInMinorDims(a, new_x, {j});
// vs[:, j] = v
vs = DynamicUpdateSliceInMinorDims(
- vs, xla::Reshape(v, ConcatVectors(batch_dims, {m, 1})), {j});
+ vs, Reshape(v, ConcatVectors(batch_dims, {m, 1})), {j});
// taus[j] = tau
taus = DynamicUpdateSliceInMinorDims(
- taus, xla::Reshape(tau, ConcatVectors(batch_dims, {1})), {j});
- return std::vector<xla::XlaOp>{a, vs, taus};
+ taus, Reshape(tau, ConcatVectors(batch_dims, {1})), {j});
+ return std::vector<XlaOp>{a, vs, taus};
};
- auto vs = xla::Zeros(builder, xla::ShapeUtil::MakeShape(
- type, ConcatVectors(batch_dims, {m, n})));
- auto taus = xla::Zeros(
- builder, xla::ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {n})));
+ auto vs = Zeros(
+ builder, ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {m, n})));
+ auto taus = Zeros(builder,
+ ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {n})));
- TF_ASSIGN_OR_RETURN(auto values,
- XlaForEachIndex(std::min(m, n), xla::S32, qr_body_fn,
- {a, vs, taus}, "qr", builder));
+ TF_ASSIGN_OR_RETURN(auto values, ForEachIndex(std::min(m, n), S32, qr_body_fn,
+ {a, vs, taus}, "qr", builder));
QRBlockResult result;
result.r = values[0];
@@ -250,24 +248,23 @@
// return W
// There is no need to return Y since at termination of the loop it is equal to
// vs.
-xla::StatusOr<xla::XlaOp> ComputeWYRepresentation(
- xla::PrimitiveType type, absl::Span<const int64> batch_dims, xla::XlaOp vs,
- xla::XlaOp taus, int64 m, int64 n,
- xla::PrecisionConfig::Precision precision) {
+StatusOr<XlaOp> ComputeWYRepresentation(PrimitiveType type,
+ absl::Span<const int64> batch_dims,
+ XlaOp vs, XlaOp taus, int64 m, int64 n,
+ PrecisionConfig::Precision precision) {
std::vector<int64> batch_dim_indices(batch_dims.size());
std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0);
int64 n_index = batch_dims.size() + 1;
- auto body_fn =
- [&](xla::XlaOp j, absl::Span<const xla::XlaOp> values,
- xla::XlaBuilder* builder) -> xla::StatusOr<std::vector<xla::XlaOp>> {
+ auto body_fn = [&](XlaOp j, absl::Span<const XlaOp> values,
+ XlaBuilder* builder) -> StatusOr<std::vector<XlaOp>> {
auto w = values[0];
auto y = values[1];
const auto vs = values[2];
const auto taus = values[3];
// Want j values in range [1, ... n).
- j = j + xla::ConstantR0<int32>(builder, 1);
+ j = j + ConstantR0<int32>(builder, 1);
// vs has shape [..., m, 1]
auto v = DynamicSliceInMinorDims(vs, {j}, {1});
// beta has shape [..., 1]
@@ -278,31 +275,31 @@
// wyv has shape [..., m, 1]
auto wyv = BatchDot(w, yv, precision);
- auto z = xla::Mul(
+ auto z = Mul(
-beta, v + wyv,
/*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index}));
w = DynamicUpdateSliceInMinorDims(w, z, {j});
y = DynamicUpdateSliceInMinorDims(y, v, {j});
- return std::vector<xla::XlaOp>{w, y, vs, taus};
+ return std::vector<XlaOp>{w, y, vs, taus};
};
- xla::XlaBuilder* builder = vs.builder();
- auto w = xla::Zeros(builder, xla::ShapeUtil::MakeShape(
- type, ConcatVectors(batch_dims, {m, n})));
+ XlaBuilder* builder = vs.builder();
+ auto w = Zeros(builder,
+ ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {m, n})));
auto y = w;
auto v = SliceInMinorDims(vs, {0}, {1});
auto beta = SliceInMinorDims(taus, {0}, {1});
y = UpdateSliceInMinorDims(y, v, {0});
- auto bv = xla::Mul(
- -beta, v,
- /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index}));
+ auto bv =
+ Mul(-beta, v,
+ /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index}));
w = UpdateSliceInMinorDims(w, bv, {0});
TF_ASSIGN_OR_RETURN(
- auto values, XlaForEachIndex(n - 1, xla::S32, body_fn, {w, y, vs, taus},
- "wy", builder));
+ auto values,
+ ForEachIndex(n - 1, S32, body_fn, {w, y, vs, taus}, "wy", builder));
return values[0];
}
@@ -323,34 +320,34 @@
// return (q, a)
// TODO(phawkins): consider using UT transformations (in the form I - V U V')
// rather than WY transformations.
-xla::StatusOr<QRDecompositionResult> QRDecomposition(
- xla::XlaOp a, bool full_matrices, int64 block_size,
- xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = a.builder();
- TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
- const int num_dims = xla::ShapeUtil::Rank(a_shape);
+StatusOr<QRDecompositionResult> QRDecomposition(
+ XlaOp a, bool full_matrices, int64 block_size,
+ PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = a.builder();
+ TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
+ const int num_dims = ShapeUtil::Rank(a_shape);
if (num_dims < 2) {
- return errors::InvalidArgument("Arguments to QR must have rank >= 2: ",
- num_dims);
+ return InvalidArgument("Arguments to QR must have rank >= 2: got shape %s",
+ a_shape.ToString());
}
- xla::PrimitiveType type = a_shape.element_type();
+ PrimitiveType type = a_shape.element_type();
- const int64 m = xla::ShapeUtil::GetDimension(a_shape, -2);
- const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
+ const int64 m = ShapeUtil::GetDimension(a_shape, -2);
+ const int64 n = ShapeUtil::GetDimension(a_shape, -1);
const int64 p = std::min(m, n);
if (block_size < 1) {
- return errors::InvalidArgument(
- "block_size argument to QR must be >= 1; got ", block_size);
+ return InvalidArgument("block_size argument to QR must be >= 1; got %d",
+ block_size);
}
const int64 num_batch_dims = num_dims - 2;
std::vector<int64> batch_dims(num_batch_dims);
for (int i = 0; i < num_batch_dims; ++i) {
- batch_dims[i] = xla::ShapeUtil::GetDimension(a_shape, i);
+ batch_dims[i] = ShapeUtil::GetDimension(a_shape, i);
}
- auto q = xla::Broadcast(xla::IdentityMatrix(builder, type, m, m), batch_dims);
+ auto q = Broadcast(IdentityMatrix(builder, type, m, m), batch_dims);
for (int64 i = 0; i < p; i += block_size) {
int64 k = std::min(block_size, p - i);
@@ -393,4 +390,4 @@
return result;
}
-} // namespace tensorflow
+} // namespace xla
diff --git a/tensorflow/compiler/tf2xla/lib/qr.h b/tensorflow/compiler/xla/client/lib/qr.h
similarity index 74%
rename from tensorflow/compiler/tf2xla/lib/qr.h
rename to tensorflow/compiler/xla/client/lib/qr.h
index 24b537a..827c8ee 100644
--- a/tensorflow/compiler/tf2xla/lib/qr.h
+++ b/tensorflow/compiler/xla/client/lib/qr.h
@@ -13,13 +13,13 @@
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_QR_H_
-#define TENSORFLOW_COMPILER_TF2XLA_LIB_QR_H_
+#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_
+#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
-namespace tensorflow {
+namespace xla {
// Computes the QR decompositions of a batch of matrices. That is,
// given a (batched) matrix a, computes an orthonormal matrix Q and an
@@ -29,14 +29,14 @@
// the block size to use.
// TODO(phawkins): handle the complex case.
struct QRDecompositionResult {
- xla::XlaOp q;
- xla::XlaOp r;
+ XlaOp q;
+ XlaOp r;
};
-xla::StatusOr<QRDecompositionResult> QRDecomposition(
- xla::XlaOp a, bool full_matrices, int64 block_size = 128,
- xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST);
+StatusOr<QRDecompositionResult> QRDecomposition(
+ XlaOp a, bool full_matrices, int64 block_size = 128,
+ PrecisionConfig::Precision precision = PrecisionConfig::HIGHEST);
-} // namespace tensorflow
+} // namespace xla
-#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_QR_H_
+#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_
diff --git a/tensorflow/compiler/xla/client/lib/qr_test.cc b/tensorflow/compiler/xla/client/lib/qr_test.cc
new file mode 100644
index 0000000..b27d364
--- /dev/null
+++ b/tensorflow/compiler/xla/client/lib/qr_test.cc
@@ -0,0 +1,93 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/client/lib/qr.h"
+
+#include "tensorflow/compiler/xla/array2d.h"
+#include "tensorflow/compiler/xla/array3d.h"
+#include "tensorflow/compiler/xla/client/lib/matrix.h"
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/compiler/xla/literal.h"
+#include "tensorflow/compiler/xla/statusor.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
+#include "tensorflow/compiler/xla/tests/literal_test_util.h"
+#include "tensorflow/compiler/xla/tests/test_macros.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace {
+
+using QrTest = xla::ClientLibraryTestBase;
+
+XLA_TEST_F(QrTest, Simple) {
+ xla::XlaBuilder builder(TestName());
+
+ xla::Array2D<float> a_vals({
+ {4, 6, 8, 10},
+ {6, 45, 54, 63},
+ {8, 54, 146, 166},
+ {10, 63, 166, 310},
+ });
+
+ xla::XlaOp a;
+ auto a_data = CreateR2Parameter<float>(a_vals, 0, "a", &builder, &a);
+ TF_ASSERT_OK_AND_ASSIGN(
+ auto result,
+ xla::QRDecomposition(a, /*full_matrices=*/true, /*block_size=*/2));
+
+ // Verifies that the decomposition composes back to the original matrix.
+ //
+ // This isn't a terribly demanding test, (e.g., we should verify that Q is
+ // orthonormal and R is upper-triangular) but it's awkward to write such tests
+ // without more linear algebra libraries. It's easier to test the numerics
+ // from Python, anyway, where we have access to numpy and scipy.
+ xla::BatchDot(result.q, result.r, xla::PrecisionConfig::HIGHEST);
+
+ ComputeAndCompareR2<float>(&builder, a_vals, {a_data.get()},
+ xla::ErrorSpec(1e-4, 1e-4));
+}
+
+XLA_TEST_F(QrTest, SimpleBatched) {
+ xla::XlaBuilder builder(TestName());
+
+ xla::Array3D<float> a_vals({
+ {
+ {4, 6, 8, 10},
+ {6, 45, 54, 63},
+ {8, 54, 146, 166},
+ {10, 63, 166, 310},
+ },
+ {
+ {16, 24, 8, 12},
+ {24, 61, 82, 48},
+ {8, 82, 456, 106},
+ {12, 48, 106, 62},
+ },
+ });
+
+ xla::XlaOp a;
+ auto a_data = CreateR3Parameter<float>(a_vals, 0, "a", &builder, &a);
+ TF_ASSERT_OK_AND_ASSIGN(
+ auto result,
+ xla::QRDecomposition(a, /*full_matrices=*/true, /*block_size=*/2));
+
+ xla::BatchDot(result.q, result.r, xla::PrecisionConfig::HIGHEST);
+
+ ComputeAndCompareR3<float>(&builder, a_vals, {a_data.get()},
+ xla::ErrorSpec(1e-4, 1e-4));
+}
+
+} // namespace
diff --git a/tensorflow/compiler/xla/client/lib/quantize.h b/tensorflow/compiler/xla/client/lib/quantize.h
new file mode 100644
index 0000000..26dbbd5
--- /dev/null
+++ b/tensorflow/compiler/xla/client/lib/quantize.h
@@ -0,0 +1,186 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QUANTIZE_H_
+#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QUANTIZE_H_
+
+#include <limits>
+#include <numeric>
+#include <vector>
+
+#include "tensorflow/compiler/xla/client/lib/constants.h"
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/compiler/xla/util.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/lib/bfloat16/bfloat16.h"
+
+namespace xla {
+
+constexpr int64 kBitsOfByte = 8;
+
+// Represents the range used for quantization
+struct QuantizedRange {
+ QuantizedRange() = default;
+ QuantizedRange(float min_in, float max_in) : min(min_in), max(max_in) {}
+
+ bool operator==(const QuantizedRange& rhs) const {
+ return this->min == rhs.min && this->max == rhs.max;
+ }
+
+ bool operator!=(const QuantizedRange& rhs) const { return !(*this == rhs); }
+
+ tensorflow::bfloat16 min = tensorflow::bfloat16(0.0f);
+ tensorflow::bfloat16 max = tensorflow::bfloat16(0.0f);
+};
+
+template <typename T>
+inline std::vector<uint32> PackToUint32(absl::Span<const T> input) {
+ const int64 kElementsPerPack = sizeof(uint32) / sizeof(T);
+ const int64 input_size = input.size();
+ const int64 output_size = CeilOfRatio(input_size, kElementsPerPack);
+
+ std::vector<uint32> output_vec;
+ constexpr int64 kShiftBits = sizeof(T) / sizeof(uint8) * kBitsOfByte;
+
+ for (int64 i = 0; i < output_size; i++) {
+ uint32 result = 0;
+ for (int64 p = 0; p < kElementsPerPack; p++) {
+ int64 index = i * kElementsPerPack + p;
+ if (index < input_size) {
+ int64 total_shift_bits = kShiftBits * (kElementsPerPack - p - 1);
+ result |= (input[index] << total_shift_bits);
+ }
+ }
+ output_vec.push_back(result);
+ }
+
+ return output_vec;
+}
+
+// Dequantize the quantized input of packed uint32 to bfloat16.
+// Only uint8 or uint16 is supported for the original unpacked input.
+// Returns a tensor of shape [d0,..., dn * unpack_size] if
+// input shape is [d0, ..., dn], where unpack_size = sizeof(unit32) / sizeof(T).
+// If transpose_output is true, will return a tensor of shape
+// [dn * unpack_size, dn-1, ..., d1, d0]. transpose_output is faster when
+// input's rank higher than 1. The input needs to be transposed to use
+// transpose_output feature.
+template <typename T>
+inline XlaOp Dequantize(XlaOp input, const QuantizedRange& range,
+ absl::string_view mode_string = "MIN_COMBINED",
+ bool transpose_output = false) {
+ XlaBuilder* const builder = input.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ float half_range =
+ !std::is_signed<T>::value
+ ? 0.0f
+ : (static_cast<float>(std::numeric_limits<T>::max()) -
+ std::numeric_limits<T>::min() + 1) /
+ 2.0f;
+ const int64 unpack_size = sizeof(uint32) / sizeof(T);
+ TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(input));
+
+ auto element_type = shape.element_type();
+ if (element_type != U32) {
+ return InvalidArgument(
+ "Only U32 is supported for input type of xla::Dequantize Op.");
+ }
+
+ // Broadcast the input to [unpack_size, d0, ..., dn] if input size is
+ // [d0, ..., dn].
+ auto broadcast_input = Broadcast(input, {unpack_size});
+
+ XlaOp iota_r1 = Iota(builder, U32, unpack_size);
+ // Highest significant bytes needs to shift more bytes than lower
+ // significant bytes.
+ XlaOp shift_bytes =
+ xla::ConstantR0<uint32>(builder, unpack_size - 1) - iota_r1;
+
+ const int bytes_of_type = sizeof(T) / sizeof(uint8);
+ std::vector<uint32> shift_vec(unpack_size, kBitsOfByte * bytes_of_type);
+ XlaOp shift_bits =
+ shift_bytes * xla::ConstantR1<uint32>(builder, shift_vec);
+
+ // Make bit_mask for different data type T.
+ uint32 bit_mask = 0x00000000;
+ for (int i = 0; i < bytes_of_type; i++) {
+ bit_mask <<= kBitsOfByte;
+ bit_mask |= 0x000000ff;
+ }
+
+ std::vector<int64> shift_transpose_dimensions(shape.dimensions_size());
+ std::iota(shift_transpose_dimensions.begin(),
+ shift_transpose_dimensions.end(), 0);
+ shift_transpose_dimensions.insert(shift_transpose_dimensions.begin(), 1,
+ shape.dimensions_size());
+
+ // Shift the input by sizeof(T) bytes and apply bit_mask to unpack.
+ XlaOp shifted_input = ShiftRightLogical(
+ broadcast_input, Transpose(Broadcast(shift_bits, shape.dimensions()),
+ shift_transpose_dimensions));
+ XlaOp unpack_input =
+ And(shifted_input, xla::ConstantR0<uint32>(builder, bit_mask));
+
+ XlaOp result;
+
+ if (mode_string == "MIN_COMBINED") {
+ const tensorflow::bfloat16 scale_factor =
+ (range.max - range.min) /
+ (static_cast<tensorflow::bfloat16>(std::numeric_limits<T>::max() -
+ std::numeric_limits<T>::min()));
+ // result = bfloat16(input + half_range) * scale_factor + range.min
+ XlaOp unpack_input_bf16 = ConvertElementType(unpack_input, BF16);
+ XlaOp half_range_bf16 = xla::ConstantR0<tensorflow::bfloat16>(
+ builder, static_cast<bfloat16>(half_range));
+ XlaOp sum = unpack_input_bf16 + half_range_bf16;
+
+ result =
+ sum * xla::ConstantR0<tensorflow::bfloat16>(builder, scale_factor) +
+ xla::ConstantR0<tensorflow::bfloat16>(builder, range.min);
+ } else {
+ // TODO(wangtao): support other modes.
+ return InvalidArgument(
+ "Only MIN_COMBINED mode is supported in xla::Dequantize Op.");
+ }
+
+ std::vector<int64> transpose_dimensions(shape.dimensions_size());
+ std::iota(transpose_dimensions.begin(), transpose_dimensions.end(), 1);
+ std::reverse(transpose_dimensions.begin(), transpose_dimensions.end());
+ transpose_dimensions.insert(transpose_dimensions.begin() + 1, 1, 0);
+
+ // Transpose the result to be [dn, unpack_size, dn-1, ..., d1, d0].
+ XlaOp transposed_result = Transpose(result, transpose_dimensions);
+
+ // Reshape to be [dn * unpack_size, dn-1, ..., d1, d0].
+ XlaOp reshaped_result = Collapse(transposed_result, {0, 1});
+
+ // Return the transpose result if transpose_output is true.
+ if (transpose_output) {
+ return reshaped_result;
+ }
+
+ // Transpose the result to be [d0, d1, ..., dn-1, dn * unpack_size].
+ std::vector<int64> result_dimensions(shape.dimensions_size());
+ std::iota(result_dimensions.begin(), result_dimensions.end(), 0);
+ std::reverse(result_dimensions.begin(), result_dimensions.end());
+
+ return Transpose(reshaped_result, result_dimensions);
+ });
+}
+
+} // namespace xla
+
+#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QUANTIZE_H_
diff --git a/tensorflow/compiler/xla/client/lib/quantize_test.cc b/tensorflow/compiler/xla/client/lib/quantize_test.cc
new file mode 100644
index 0000000..be3603d
--- /dev/null
+++ b/tensorflow/compiler/xla/client/lib/quantize_test.cc
@@ -0,0 +1,337 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/client/lib/quantize.h"
+
+#include <limits>
+
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
+#include "tensorflow/compiler/xla/tests/test_macros.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/compiler/xla/util.h"
+
+namespace xla {
+namespace {
+
+using bfloat16 = tensorflow::bfloat16;
+
+template <typename NativeT>
+std::vector<NativeT> GenerateInput() {
+ std::vector<NativeT> input;
+
+ for (int64 i = std::numeric_limits<NativeT>::min();
+ i < std::numeric_limits<NativeT>::max(); ++i) {
+ input.push_back(static_cast<NativeT>(i));
+ }
+
+ return input;
+}
+
+template <typename NativeT>
+Array2D<NativeT> GenerateLargeSizeInput(int num_columns, int num_rows) {
+ Array2D<NativeT> input(num_columns, num_rows);
+
+ input.FillRandom(6, 128);
+
+ return input;
+}
+
+template <typename NativeT>
+Array2D<uint32> PackLargeInput(Array2D<NativeT> &input) {
+ const int64 size_per_pack = sizeof(uint32) / sizeof(NativeT);
+ int64 width = input.width();
+
+ int64 padded_output_width = CeilOfRatio(width, size_per_pack);
+
+ Array2D<uint32> pack_input(input.height(), padded_output_width);
+
+ for (int h = 0; h < input.height(); h++) {
+ std::vector<NativeT> input_row;
+ for (int w = 0; w < width; w++) {
+ input_row.push_back(input({h, w}));
+ }
+
+ auto pack_input_vec = PackToUint32<uint8>(input_row);
+
+ for (int w = 0; w < padded_output_width; w++) {
+ pack_input(h, w) = pack_input_vec[w];
+ }
+ }
+
+ return pack_input;
+}
+
+template <typename NativeT>
+Array2D<bfloat16> GenerateLargeSizeMinCombinedOutput(
+ Array2D<NativeT> &input, const QuantizedRange &range,
+ bool transpose_output = false) {
+ const int64 size_per_pack = sizeof(uint32) / sizeof(NativeT);
+ int64 width = input.width();
+
+ int64 padded_output_width = CeilOfRatio(width, size_per_pack) * size_per_pack;
+
+ int64 output_height;
+ int64 output_width;
+
+ if (transpose_output) {
+ output_height = padded_output_width;
+ output_width = input.height();
+ } else {
+ output_height = input.height();
+ output_width = padded_output_width;
+ }
+
+ Array2D<bfloat16> output(output_height, output_width, bfloat16(0.0));
+
+ float half_range =
+ !std::is_signed<NativeT>::value
+ ? 0.0f
+ : (static_cast<float>(std::numeric_limits<NativeT>::max() -
+ std::numeric_limits<NativeT>::min() + 1)) /
+ 2.0f;
+ const bfloat16 scale_factor =
+ (range.max - range.min) /
+ (static_cast<bfloat16>(std::numeric_limits<NativeT>::max() -
+ std::numeric_limits<NativeT>::min()));
+
+ for (int h = 0; h < input.height(); h++) {
+ std::vector<NativeT> input_row;
+ for (int w = 0; w < width; w++) {
+ bfloat16 result =
+ static_cast<bfloat16>(input(h, w) + half_range) * scale_factor +
+ range.min;
+ if (transpose_output) {
+ output(w, h) = result;
+ } else {
+ output(h, w) = result;
+ }
+ }
+ }
+
+ return output;
+}
+
+template <typename NativeT>
+std::vector<bfloat16> GenerateMinCombinedOutput(const QuantizedRange &range) {
+ float half_range =
+ !std::is_signed<NativeT>::value
+ ? 0.0f
+ : (static_cast<float>(std::numeric_limits<NativeT>::max() -
+ std::numeric_limits<NativeT>::min() + 1)) /
+ 2.0f;
+ const bfloat16 scale_factor =
+ (range.max - range.min) /
+ (static_cast<bfloat16>(std::numeric_limits<NativeT>::max() -
+ std::numeric_limits<NativeT>::min()));
+ std::vector<bfloat16> output;
+ for (int64 i = std::numeric_limits<NativeT>::min();
+ i < std::numeric_limits<NativeT>::max(); ++i) {
+ bfloat16 result =
+ static_cast<bfloat16>(i + half_range) * scale_factor + range.min;
+ output.push_back(result);
+ }
+
+ const int64 pack_size = sizeof(uint32) / sizeof(NativeT);
+ const int64 output_size = output.size();
+
+ int64 num_tailing_zeros =
+ CeilOfRatio(output_size, pack_size) * pack_size - output_size;
+
+ output.insert(output.end(), num_tailing_zeros, bfloat16(0.0));
+ return output;
+}
+
+// TODO(wangtao): add a test to make sure this op is the inverse of the existing
+// TF quantize op defined in: third_party/tensorflow/core/kernels/quantize_op.cc
+
+using DequantizeTest = ClientLibraryTestBase;
+
+TEST(PackTest, PackUint8ToUint32) {
+ std::vector<uint8> input = {0xAB, 0x0B, 0x00, 0xF0, 0x01};
+ auto output = PackToUint32<uint8>(input);
+ EXPECT_THAT(output, ::testing::ElementsAre(0xAB0B00F0, 0x01000000));
+}
+
+TEST(PackTest, PackInt8ToUint32) {
+ std::vector<int8> input = {static_cast<signed char>(0x81), 0x0B, 0x00, 0x20,
+ 0x01};
+ auto output = PackToUint32<int8>(input);
+ EXPECT_THAT(output, ::testing::ElementsAre(0x810B0020, 0x01000000));
+}
+
+TEST(PackTest, PackUint8ToUint32PerfectSize) {
+ std::vector<uint8> input = {3, 2, 1, 0};
+ auto output = PackToUint32<uint8>(input);
+ EXPECT_THAT(output, ::testing::ElementsAre(0x03020100));
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint16R1) {
+ XlaBuilder builder(TestName());
+ auto input = GenerateInput<uint16>();
+ auto x = ConstantR1<uint32>(&builder, PackToUint32<uint16>(input));
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint16>(x, range, "MIN_COMBINED");
+ auto expected = GenerateMinCombinedOutput<uint16>(range);
+ ComputeAndCompareR1<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8R1) {
+ XlaBuilder builder(TestName());
+ auto input = GenerateInput<uint8>();
+ auto x = ConstantR1<uint32>(&builder, PackToUint32<uint8>(input));
+ QuantizedRange range(0, 127.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
+ auto expected = GenerateMinCombinedOutput<uint8>(range);
+ ComputeAndCompareR1<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8R2) {
+ XlaBuilder builder(TestName());
+ std::vector<std::vector<uint8>> input = {
+ {0, 1, 2, 3},
+ {4, 5, 6, 7},
+ {8, 9, 10, 11},
+ {12, 13, 16, 15},
+ };
+ auto x = ConstantR2<uint32>(&builder, {{PackToUint32<uint8>(input[0])[0]},
+ {PackToUint32<uint8>(input[1])[0]},
+ {PackToUint32<uint8>(input[2])[0]},
+ {PackToUint32<uint8>(input[3])[0]}});
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
+ const Array2D<bfloat16> expected = {
+ {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0)},
+ {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0)},
+ {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0)},
+ {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0)},
+ };
+ ComputeAndCompareR2<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TransposeOutput) {
+ XlaBuilder builder(TestName());
+ std::vector<std::vector<uint8>> input = {
+ {0, 1, 2, 3},
+ {4, 5, 6, 7},
+ {8, 9, 10, 11},
+ {12, 13, 16, 15},
+ };
+ auto x = ConstantR2<uint32>(&builder, {{PackToUint32<uint8>(input[0])[0]},
+ {PackToUint32<uint8>(input[1])[0]},
+ {PackToUint32<uint8>(input[2])[0]},
+ {PackToUint32<uint8>(input[3])[0]}});
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
+ const Array2D<bfloat16> expected = {
+ {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)},
+ {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)},
+ {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)},
+ {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)},
+ };
+ ComputeAndCompareR2<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZero) {
+ XlaBuilder builder(TestName());
+ std::vector<std::vector<uint8>> input = {
+ {0, 1, 2, 3, 16},
+ {4, 5, 6, 7, 17},
+ {8, 9, 10, 11, 18},
+ {12, 13, 16, 15, 19},
+ };
+ auto x = ConstantR2<uint32>(
+ &builder,
+ {{PackToUint32<uint8>(input[0])[0], PackToUint32<uint8>(input[0])[1]},
+ {PackToUint32<uint8>(input[1])[0], PackToUint32<uint8>(input[1])[1]},
+ {PackToUint32<uint8>(input[2])[0], PackToUint32<uint8>(input[2])[1]},
+ {PackToUint32<uint8>(input[3])[0], PackToUint32<uint8>(input[3])[1]}});
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
+
+ const Array2D<bfloat16> expected = {
+ {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0),
+ bfloat16(16.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0),
+ bfloat16(17.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0),
+ bfloat16(18.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0),
+ bfloat16(19.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ };
+ ComputeAndCompareR2<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZeroTransposeOutput) {
+ XlaBuilder builder(TestName());
+ std::vector<std::vector<uint8>> input = {
+ {0, 1, 2, 3, 16},
+ {4, 5, 6, 7, 17},
+ {8, 9, 10, 11, 18},
+ {12, 13, 16, 15, 19},
+ };
+ auto x = ConstantR2<uint32>(
+ &builder,
+ {{PackToUint32<uint8>(input[0])[0], PackToUint32<uint8>(input[0])[1]},
+ {PackToUint32<uint8>(input[1])[0], PackToUint32<uint8>(input[1])[1]},
+ {PackToUint32<uint8>(input[2])[0], PackToUint32<uint8>(input[2])[1]},
+ {PackToUint32<uint8>(input[3])[0], PackToUint32<uint8>(input[3])[1]}});
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
+
+ const Array2D<bfloat16> expected = {
+ {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)},
+ {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)},
+ {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)},
+ {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)},
+ {bfloat16(16.0), bfloat16(17.0), bfloat16(18.0), bfloat16(19.0)},
+ {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)},
+ };
+ ComputeAndCompareR2<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTest) {
+ XlaBuilder builder(TestName());
+ Array2D<uint8> input = GenerateLargeSizeInput<uint8>(500, 3547);
+ Array2D<uint32> input_packed = PackLargeInput<uint8>(input);
+
+ auto x = ConstantR2FromArray2D<uint32>(&builder, input_packed);
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED");
+
+ const Array2D<bfloat16> expected =
+ GenerateLargeSizeMinCombinedOutput<uint8>(input, range);
+ ComputeAndCompareR2<bfloat16>(&builder, expected, {});
+}
+
+XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTestTransposeOutput) {
+ XlaBuilder builder(TestName());
+ Array2D<uint8> input = GenerateLargeSizeInput<uint8>(500, 3547);
+ Array2D<uint32> input_packed = PackLargeInput<uint8>(input);
+
+ auto x = ConstantR2FromArray2D<uint32>(&builder, input_packed);
+ QuantizedRange range(0, 255.0f);
+ xla::Dequantize<uint8>(x, range, "MIN_COMBINED", /*transpose_output=*/true);
+
+ const Array2D<bfloat16> expected = GenerateLargeSizeMinCombinedOutput<uint8>(
+ input, range, /*transpose_output=*/true);
+ ComputeAndCompareR2<bfloat16>(&builder, expected, {});
+}
+
+} // namespace
+} // namespace xla
diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc
index a95bbf2..5db9d10 100644
--- a/tensorflow/compiler/xla/client/lib/testing.cc
+++ b/tensorflow/compiler/xla/client/lib/testing.cc
@@ -59,22 +59,25 @@
return Tuple(builder, parts);
}
-std::unique_ptr<GlobalData> MakeFakeDataViaDeviceOrDie(const Shape& shape,
- Client* client) {
+std::unique_ptr<GlobalData> MakeFakeDataViaDeviceOrDie(
+ const Shape& shape, Client* client, DebugOptions* debug_opts) {
XlaBuilder b(absl::StrCat("make_fake_", ShapeUtil::HumanString(shape)));
BuildFakeDataOpOnDevice(shape, &b);
XlaComputation computation = b.Build().ConsumeValueOrDie();
auto execution_options = CreateDefaultExecutionOptions();
*execution_options.mutable_shape_with_output_layout() = shape.ToProto();
+ if (debug_opts) {
+ *execution_options.mutable_debug_options() = *debug_opts;
+ }
return client->Execute(computation, /*arguments=*/{}, &execution_options)
.ConsumeValueOrDie();
}
} // namespace
-std::unique_ptr<GlobalData> MakeFakeDataOrDie(const Shape& shape,
- Client* client) {
+std::unique_ptr<GlobalData> MakeFakeDataOrDie(
+ const Shape& shape, Client* client, DebugOptions* debug_opts /*=nullptr*/) {
if (DataSizeOfShape(shape) < (1LL << 20)) {
StatusOr<Literal> literal_status = MakeFakeLiteral(shape);
if (!literal_status.ok()) {
@@ -82,24 +85,25 @@
// an on-device computation.
CHECK_EQ(literal_status.status().code(),
tensorflow::error::UNIMPLEMENTED);
- return MakeFakeDataViaDeviceOrDie(shape, client);
+ return MakeFakeDataViaDeviceOrDie(shape, client, debug_opts);
}
return client->TransferToServer(literal_status.ValueOrDie()).ValueOrDie();
}
// If the data is large, generate it on-device.
- return MakeFakeDataViaDeviceOrDie(shape, client);
+ return MakeFakeDataViaDeviceOrDie(shape, client, debug_opts);
}
std::vector<std::unique_ptr<GlobalData>> MakeFakeArgumentsOrDie(
- const XlaComputation& computation, Client* client) {
+ const XlaComputation& computation, Client* client,
+ DebugOptions* debug_opts /*=nullptr*/) {
CHECK(computation.proto().has_host_program_shape())
<< "Computation should have progran shape.";
auto program_shape = computation.proto().host_program_shape();
std::vector<std::unique_ptr<GlobalData>> results;
for (const ShapeProto& shape : program_shape.parameters()) {
- results.push_back(MakeFakeDataOrDie(Shape(shape), client));
+ results.push_back(MakeFakeDataOrDie(Shape(shape), client, debug_opts));
}
return results;
}
diff --git a/tensorflow/compiler/xla/client/lib/testing.h b/tensorflow/compiler/xla/client/lib/testing.h
index 03695ce..428fa3e 100644
--- a/tensorflow/compiler/xla/client/lib/testing.h
+++ b/tensorflow/compiler/xla/client/lib/testing.h
@@ -29,14 +29,19 @@
// Generates fake data of the given shape on the device or dies. The fake data
// is created by performing a computation on the device rather than transferring
// data from the host to the device.
-std::unique_ptr<GlobalData> MakeFakeDataOrDie(const Shape& shape,
- Client* client);
+//
+// The optional DebugOptions are used when generating fake data on the device.
+std::unique_ptr<GlobalData> MakeFakeDataOrDie(
+ const Shape& shape, Client* client, DebugOptions* debug_opts = nullptr);
// Returns vector of GlobalData handles of fake data (created using
// MakeFakeDataOrDie) that are correctly shaped arguments for the given
// xla computation.
+//
+// The optional DebugOptions are used when generating fake data on the device.
std::vector<std::unique_ptr<GlobalData>> MakeFakeArgumentsOrDie(
- const XlaComputation& computation, Client* client);
+ const XlaComputation& computation, Client* client,
+ DebugOptions* debug_opts = nullptr);
} // namespace xla
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc b/tensorflow/compiler/xla/client/lib/triangular_solve.cc
similarity index 67%
rename from tensorflow/compiler/tf2xla/lib/triangular_solve.cc
rename to tensorflow/compiler/xla/client/lib/triangular_solve.cc
index 192a61d..c5a1d34 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
+++ b/tensorflow/compiler/xla/client/lib/triangular_solve.cc
@@ -13,12 +13,11 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h"
+#include "tensorflow/compiler/xla/client/lib/triangular_solve.h"
#include <memory>
#include <vector>
-#include "tensorflow/compiler/tf2xla/lib/util.h"
#include "tensorflow/compiler/xla/client/lib/constants.h"
#include "tensorflow/compiler/xla/client/lib/math.h"
#include "tensorflow/compiler/xla/client/lib/matrix.h"
@@ -30,21 +29,20 @@
#include "tensorflow/compiler/xla/status_macros.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/util.h"
-#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/math/math_util.h"
-namespace tensorflow {
+namespace xla {
// Get the diagonal blocks of the coefficient matrix
-xla::XlaOp DiagonalBlocks(xla::XlaOp a, int64 block_size) {
- xla::XlaBuilder* builder = a.builder();
- return builder->ReportErrorOrReturn([&]() -> xla::StatusOr<xla::XlaOp> {
- TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(a));
- int ndims = xla::ShapeUtil::Rank(shape);
- int64 n = xla::ShapeUtil::GetDimension(shape, -1);
+XlaOp DiagonalBlocks(XlaOp a, int64 block_size) {
+ XlaBuilder* builder = a.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(a));
+ int ndims = ShapeUtil::Rank(shape);
+ int64 n = ShapeUtil::GetDimension(shape, -1);
int64 num_blocks = n / block_size;
- xla::XlaOp diag_blocks;
+ XlaOp diag_blocks;
// If the coefficient matrix is exactly the block size, we just add a
// singleton dimension i.e. [..., n, n] -> [..., 1, n, n]
@@ -59,13 +57,13 @@
if (n > block_size) {
// Construct the starting indices of the diagonal blocks
auto start_indices =
- Transpose(Broadcast(Mul(Iota(builder, xla::S32, num_blocks),
- xla::ConstantR0<int32>(builder, block_size)),
+ Transpose(Broadcast(Mul(Iota(builder, S32, num_blocks),
+ ConstantR0<int32>(builder, block_size)),
/*broadcast_sizes=*/{2}),
/*permutation=*/{1, 0});
// Gather the diagonal blocks
- xla::GatherDimensionNumbers dim_numbers;
+ GatherDimensionNumbers dim_numbers;
dim_numbers.add_offset_dims(ndims - 1);
dim_numbers.add_offset_dims(ndims);
dim_numbers.add_start_index_map(ndims - 2);
@@ -81,7 +79,7 @@
// Pad with zeros
auto last_blocks =
SliceInMinorDims(a, {n - n % block_size, n - n % block_size}, {n, n});
- xla::PaddingConfig config = xla::MakeNoPaddingConfig(ndims);
+ PaddingConfig config = MakeNoPaddingConfig(ndims);
int64 padding = block_size - n % block_size;
config.mutable_dimensions(ndims - 1)->set_edge_padding_high(padding);
config.mutable_dimensions(ndims - 2)->set_edge_padding_high(padding);
@@ -90,9 +88,8 @@
// Add a singleton dimension
// i.e. [..., block_size, block_size] -> [..., 1, block_size, block_size]
- TF_ASSIGN_OR_RETURN(xla::Shape blocks_shape,
- builder->GetShape(last_blocks));
- auto shape_dims = xla::AsInt64Slice(blocks_shape.dimensions());
+ TF_ASSIGN_OR_RETURN(Shape blocks_shape, builder->GetShape(last_blocks));
+ auto shape_dims = AsInt64Slice(blocks_shape.dimensions());
auto last_blocks_dims = std::vector<int64>(ndims);
std::copy(shape_dims.begin(), shape_dims.end(), last_blocks_dims.begin());
last_blocks_dims.insert(last_blocks_dims.end() - 2, 1);
@@ -101,7 +98,7 @@
// Concatenate with the other blocks if necessary
if (n > block_size) {
diag_blocks =
- xla::ConcatInDim(builder, {diag_blocks, last_blocks}, ndims - 2);
+ ConcatInDim(builder, {diag_blocks, last_blocks}, ndims - 2);
} else {
diag_blocks = last_blocks;
}
@@ -111,16 +108,16 @@
});
}
-xla::XlaOp InvertDiagonalBlocks(xla::XlaOp diag_blocks, bool lower,
- bool transpose_a, bool conjugate_a,
- xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = diag_blocks.builder();
- return builder->ReportErrorOrReturn([&]() -> xla::StatusOr<xla::XlaOp> {
+XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a,
+ bool conjugate_a,
+ PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = diag_blocks.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
// Input is a batch of square lower triangular square matrices. Its shape is
// (..., size, size). We resize this to (num_blocks, size, size).
- TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(diag_blocks));
- int64 block_size = xla::ShapeUtil::GetDimension(shape, -1);
- int64 num_blocks = xla::ShapeUtil::ElementsIn(shape) /
+ TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(diag_blocks));
+ int64 block_size = ShapeUtil::GetDimension(shape, -1);
+ int64 num_blocks = ShapeUtil::ElementsIn(shape) /
tensorflow::MathUtil::IPow(block_size, 2);
diag_blocks = Reshape(diag_blocks, {num_blocks, block_size, block_size});
@@ -132,9 +129,9 @@
// zero (which can happen if the last block was padded) otherwise it will
// introduce nans which will propagate
auto diags = GetMatrixDiagonal(diag_blocks);
- TF_ASSIGN_OR_RETURN(xla::Shape diags_shape, builder->GetShape(diags));
+ TF_ASSIGN_OR_RETURN(Shape diags_shape, builder->GetShape(diags));
auto one = ScalarLike(diags, 1);
- auto ones = Broadcast(one, xla::AsInt64Slice(diags_shape.dimensions()));
+ auto ones = Broadcast(one, AsInt64Slice(diags_shape.dimensions()));
diags = Select(Eq(diags, Zero(builder, shape.element_type())), ones, diags);
auto scaled_diag_blocks = Div(diag_blocks, diags, {0, 2});
@@ -160,40 +157,40 @@
auto start_index = (lower) ? 0 : block_size - 1;
auto output_block = DynamicUpdateSlice(
neg_identity, pos_one,
- /*start_indices=*/xla::ConstantR1<int>(builder, 2, start_index));
+ /*start_indices=*/ConstantR1<int>(builder, 2, start_index));
// Broadcast diag([1, -1, -1, ...]) to every block
- xla::XlaOp output = Broadcast(output_block,
- /*broadcast_sizes=*/{num_blocks});
+ XlaOp output = Broadcast(output_block,
+ /*broadcast_sizes=*/{num_blocks});
// Now we construct a loop that performs matrix-vector multiplications
// inverting the blocks one row at a time
- std::vector<xla::Shape> tuple_shapes = {
+ std::vector<Shape> tuple_shapes = {
// The loop iteration counter is a scalar, incremented each iteration.
- xla::ShapeUtil::MakeShape(xla::S32, {}),
+ ShapeUtil::MakeShape(S32, {}),
// The output has the shape of A, with one row updated each iteration.
- xla::ShapeUtil::MakeShape(shape.element_type(),
- {num_blocks, block_size, block_size}),
+ ShapeUtil::MakeShape(shape.element_type(),
+ {num_blocks, block_size, block_size}),
// The input is a loop invariant.
- xla::ShapeUtil::MakeShape(shape.element_type(),
- {num_blocks, block_size, block_size})};
- xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes);
+ ShapeUtil::MakeShape(shape.element_type(),
+ {num_blocks, block_size, block_size})};
+ Shape tuple_shape = ShapeUtil::MakeTupleShape(tuple_shapes);
- auto init_i = One(builder, xla::S32);
- auto init = xla::Tuple(builder, {init_i, output, scaled_diag_blocks});
+ auto init_i = One(builder, S32);
+ auto init = Tuple(builder, {init_i, output, scaled_diag_blocks});
// Construct the loop condition function.
- std::unique_ptr<xla::XlaBuilder> condb =
+ std::unique_ptr<XlaBuilder> condb =
builder->CreateSubBuilder("InvertDiagCond");
{
auto i = GetTupleElement(
Parameter(condb.get(), 0, tuple_shape, "InvertDiagCondTuple"), 0);
- Lt(i, xla::ConstantR0<int32>(condb.get(), block_size));
+ Lt(i, ConstantR0<int32>(condb.get(), block_size));
}
TF_ASSIGN_OR_RETURN(auto cond, condb->Build());
// Construct the loop body function.
- std::unique_ptr<xla::XlaBuilder> bodyb =
+ std::unique_ptr<XlaBuilder> bodyb =
builder->CreateSubBuilder("InvertDiagBody");
{
auto input_tuple =
@@ -203,21 +200,21 @@
auto body_out = GetTupleElement(input_tuple, 1);
auto body_input = GetTupleElement(input_tuple, 2);
- auto zero = xla::ConstantR1<int32>(bodyb.get(), 1, 0);
+ auto zero = ConstantR1<int32>(bodyb.get(), 1, 0);
auto j = (lower) ? i : ScalarLike(i, block_size - 1) - i;
auto start_indices =
- xla::ConcatInDim(bodyb.get(), {zero, Reshape(j, {1}), zero}, 0);
+ ConcatInDim(bodyb.get(), {zero, Reshape(j, {1}), zero}, 0);
auto input_row =
DynamicSlice(body_input, start_indices,
/*slice_sizes=*/{num_blocks, 1, block_size});
// We want -L21 L11^{-1}
- xla::DotDimensionNumbers dnums;
+ DotDimensionNumbers dnums;
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
- xla::PrecisionConfig precision_proto;
+ PrecisionConfig precision_proto;
precision_proto.add_operand_precision(precision);
precision_proto.add_operand_precision(precision);
auto update = -DotGeneral(input_row, body_out, dnums, &precision_proto);
@@ -225,7 +222,7 @@
body_out = DynamicUpdateSlice(body_out, update, start_indices);
auto next_i = i + ScalarLike(i, 1);
- xla::Tuple(bodyb.get(), {next_i, body_out, body_input});
+ Tuple(bodyb.get(), {next_i, body_out, body_input});
}
TF_ASSIGN_OR_RETURN(auto body, bodyb->Build());
@@ -239,27 +236,26 @@
/*broadcast_dimensions=*/{0, 1});
// Reshape back to original batch major dimensions
- return Reshape(inv_diag_blocks, xla::AsInt64Slice(shape.dimensions()));
+ return Reshape(inv_diag_blocks, AsInt64Slice(shape.dimensions()));
});
}
-xla::XlaOp SolveWithInvertedDiagonalBlocks(
- xla::XlaOp a, xla::XlaOp b, xla::XlaOp inv_diag_blocks, bool left_side,
- bool lower, bool transpose_a, bool conjugate_a,
- xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = a.builder();
- return builder->ReportErrorOrReturn([&]() -> xla::StatusOr<xla::XlaOp> {
- TF_ASSIGN_OR_RETURN(xla::Shape blocks_shape,
- builder->GetShape(inv_diag_blocks));
- TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b));
- int64 block_size = xla::ShapeUtil::GetDimension(blocks_shape, -1);
+XlaOp SolveWithInvertedDiagonalBlocks(XlaOp a, XlaOp b, XlaOp inv_diag_blocks,
+ bool left_side, bool lower,
+ bool transpose_a, bool conjugate_a,
+ PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = a.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ TF_ASSIGN_OR_RETURN(Shape blocks_shape, builder->GetShape(inv_diag_blocks));
+ TF_ASSIGN_OR_RETURN(Shape b_shape, builder->GetShape(b));
+ int64 block_size = ShapeUtil::GetDimension(blocks_shape, -1);
- TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
- int64 ndims = xla::ShapeUtil::Rank(a_shape);
- int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
+ TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
+ int64 ndims = ShapeUtil::Rank(a_shape);
+ int64 n = ShapeUtil::GetDimension(a_shape, -1);
int64 num_blocks = n / block_size + (n % block_size != 0);
int64 m_dim = (left_side) ? -1 : -2;
- int64 m = xla::ShapeUtil::GetDimension(b_shape, m_dim);
+ int64 m = ShapeUtil::GetDimension(b_shape, m_dim);
// Initialize the solution
auto x = ZerosLike(b);
@@ -295,7 +291,7 @@
}
auto b_row = SliceInMinorDims(b, start, end);
- xla::XlaOp remainder;
+ XlaOp remainder;
if (i == 0) {
remainder = b_row;
} else {
@@ -322,11 +318,10 @@
}
}
- xla::XlaOp x_update;
- auto zero = Zero(builder, xla::S32);
- auto start_index =
- xla::ConstantR0WithType(builder, xla::S32, j * block_size);
- std::vector<xla::XlaOp> update_starts = {start_index, zero};
+ XlaOp x_update;
+ auto zero = Zero(builder, S32);
+ auto start_index = ConstantR0WithType(builder, S32, j * block_size);
+ std::vector<XlaOp> update_starts = {start_index, zero};
if (left_side) {
x_update = BatchDot(MaybeTransposeInMinorDims(inv_block, transpose_a),
remainder, precision);
@@ -343,24 +338,24 @@
});
}
-xla::XlaOp TriangularSolve(xla::XlaOp a, xla::XlaOp b, bool left_side,
- bool lower, bool transpose_a, bool conjugate_a,
- int64 block_size,
- xla::PrecisionConfig::Precision precision) {
- xla::XlaBuilder* builder = a.builder();
- return builder->ReportErrorOrReturn([&]() -> xla::StatusOr<xla::XlaOp> {
- TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
- TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b));
- if (xla::ShapeUtil::Rank(a_shape) != xla::ShapeUtil::Rank(b_shape)) {
- return errors::InvalidArgument(
- "Arguments to TriangularSolve have different ranks: ",
- xla::ShapeUtil::HumanString(a_shape), " vs. ",
- xla::ShapeUtil::HumanString(b_shape));
+XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower,
+ bool transpose_a, bool conjugate_a, int64 block_size,
+ PrecisionConfig::Precision precision) {
+ XlaBuilder* builder = a.builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a));
+ TF_ASSIGN_OR_RETURN(Shape b_shape, builder->GetShape(b));
+ if (ShapeUtil::Rank(a_shape) != ShapeUtil::Rank(b_shape)) {
+ return InvalidArgument(
+ "Arguments to TriangularSolve have shapes with different ranks: "
+ "%s vs. %s",
+ ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape));
}
- const int64 ndims = xla::ShapeUtil::Rank(a_shape);
+ const int64 ndims = ShapeUtil::Rank(a_shape);
if (ndims < 2) {
- return errors::InvalidArgument(
- "Arguments to TriangularSolve must have rank >= 2: ", ndims);
+ return InvalidArgument(
+ "Arguments to TriangularSolve was rank %d but must have rank >= 2.",
+ ndims);
}
// The batch dimensions must be equal.
std::vector<int64> batch_dimensions;
@@ -368,32 +363,33 @@
int64 a_size = a_shape.dimensions(i);
int64 b_size = b_shape.dimensions(i);
if (a_size != b_size) {
- return errors::InvalidArgument(
- "Batch dimensions of arguments to TriangularSolve must be equal: ",
- xla::ShapeUtil::HumanString(a_shape), " vs ",
- xla::ShapeUtil::HumanString(b_shape));
+ return InvalidArgument(
+ "Batch dimensions of arguments to TriangularSolve must be equal; "
+ "shapes were %s and %s.",
+ ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape));
}
batch_dimensions.push_back(a_size);
}
- if (xla::ShapeUtil::GetDimension(a_shape, -1) !=
- xla::ShapeUtil::GetDimension(a_shape, -2)) {
- return errors::InvalidArgument(
- "The 'a' arguments to TriangularSolve must be square matrices: ",
- xla::ShapeUtil::HumanString(a_shape));
+ if (ShapeUtil::GetDimension(a_shape, -1) !=
+ ShapeUtil::GetDimension(a_shape, -2)) {
+ return InvalidArgument(
+ "The 'a' argument to TriangularSolve must be a batched square matrix;"
+ " shape was: %s",
+ ShapeUtil::HumanString(a_shape));
}
- const int64 m = xla::ShapeUtil::GetDimension(b_shape, -2);
- const int64 n = xla::ShapeUtil::GetDimension(b_shape, -1);
- if ((left_side ? m : n) != xla::ShapeUtil::GetDimension(a_shape, -1)) {
- return errors::InvalidArgument(
- "Arguments to TriangularSolve have incompatible matrix shapes: ",
- xla::ShapeUtil::HumanString(a_shape), " vs ",
- xla::ShapeUtil::HumanString(b_shape));
+ const int64 m = ShapeUtil::GetDimension(b_shape, -2);
+ const int64 n = ShapeUtil::GetDimension(b_shape, -1);
+ if ((left_side ? m : n) != ShapeUtil::GetDimension(a_shape, -1)) {
+ return InvalidArgument(
+ "Arguments to TriangularSolve have incompatible matrix shapes %s and "
+ "%s",
+ ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape));
}
if (block_size < 1) {
- return errors::InvalidArgument(
- "block_size argument to TriangularSolve must be >= 1; got ",
+ return InvalidArgument(
+ "block_size argument to TriangularSolve must be >= 1; got %d",
block_size);
}
@@ -413,4 +409,4 @@
});
}
-} // namespace tensorflow
+} // namespace xla
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.h b/tensorflow/compiler/xla/client/lib/triangular_solve.h
similarity index 87%
rename from tensorflow/compiler/tf2xla/lib/triangular_solve.h
rename to tensorflow/compiler/xla/client/lib/triangular_solve.h
index 2303234..50a3b30 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve.h
+++ b/tensorflow/compiler/xla/client/lib/triangular_solve.h
@@ -13,13 +13,13 @@
limitations under the License.
==============================================================================*/
-#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_
-#define TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_
+#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_
+#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
-namespace tensorflow {
+namespace xla {
// Solves systems of linear equations with lower or upper triangular coefficient
// matrices by forward- or back-substitution. Broadcasting along leading
@@ -57,11 +57,11 @@
//
// Uses a blocked algorithm if `block_size` is > 1; if block_size == 1 then no
// blocking is used.
-xla::XlaOp TriangularSolve(
- xla::XlaOp a, xla::XlaOp b, bool left_side, bool lower, bool transpose_a,
+XlaOp TriangularSolve(
+ XlaOp a, XlaOp b, bool left_side, bool lower, bool transpose_a,
bool conjugate_a, int64 block_size = 128,
- xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST);
+ PrecisionConfig::Precision precision = PrecisionConfig::HIGHEST);
-} // namespace tensorflow
+} // namespace xla
-#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_
+#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc b/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc
similarity index 98%
rename from tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc
rename to tensorflow/compiler/xla/client/lib/triangular_solve_test.cc
index aeebf16..f6a70d6 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc
+++ b/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc
@@ -13,7 +13,7 @@
limitations under the License.
==============================================================================*/
-#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h"
+#include "tensorflow/compiler/xla/client/lib/triangular_solve.h"
#include <memory>
#include <numeric>
@@ -30,7 +30,7 @@
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/core/lib/core/status_test_util.h"
-namespace tensorflow {
+namespace xla {
namespace {
using TriangularSolveTest = xla::ClientLibraryTestBase;
@@ -330,4 +330,4 @@
}
} // namespace
-} // namespace tensorflow
+} // namespace xla
diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc
index aaa5d69..049cd15 100644
--- a/tensorflow/compiler/xla/client/local_client.cc
+++ b/tensorflow/compiler/xla/client/local_client.cc
@@ -71,9 +71,9 @@
"parameter "
"%d: want %s, got %s",
i,
- ShapeUtil::HumanString(
+ ShapeUtil::HumanStringWithLayout(
computation_layout.parameter_layout(i).shape()),
- ShapeUtil::HumanString(arguments[i]->on_host_shape()));
+ ShapeUtil::HumanStringWithLayout(arguments[i]->on_host_shape()));
}
}
diff --git a/tensorflow/compiler/xla/debug_options_flags.cc b/tensorflow/compiler/xla/debug_options_flags.cc
index d7e7b9e..a9a9164 100644
--- a/tensorflow/compiler/xla/debug_options_flags.cc
+++ b/tensorflow/compiler/xla/debug_options_flags.cc
@@ -22,49 +22,49 @@
#include "tensorflow/compiler/xla/parse_flags_from_env.h"
namespace xla {
-namespace {
-DebugOptions* flag_values;
-std::vector<tensorflow::Flag>* flag_objects;
-std::once_flag flags_init;
-
-void SetDebugOptionsDefaults(DebugOptions* flags) {
- flags->set_xla_llvm_enable_alias_scope_metadata(true);
- flags->set_xla_llvm_enable_noalias_metadata(true);
- flags->set_xla_llvm_enable_invariant_load_metadata(true);
- flags->set_xla_llvm_disable_expensive_passes(false);
- flags->set_xla_backend_optimization_level(3);
- flags->set_xla_cpu_multi_thread_eigen(true);
- flags->set_xla_gpu_cuda_data_dir("./cuda_sdk_lib");
- flags->set_xla_eliminate_hlo_implicit_broadcast(true);
+DebugOptions DefaultDebugOptionsIgnoringFlags() {
+ DebugOptions opts;
+ opts.set_xla_llvm_enable_alias_scope_metadata(true);
+ opts.set_xla_llvm_enable_noalias_metadata(true);
+ opts.set_xla_llvm_enable_invariant_load_metadata(true);
+ opts.set_xla_llvm_disable_expensive_passes(false);
+ opts.set_xla_backend_optimization_level(3);
+ opts.set_xla_cpu_multi_thread_eigen(true);
+ opts.set_xla_gpu_cuda_data_dir("./cuda_sdk_lib");
+ opts.set_xla_eliminate_hlo_implicit_broadcast(true);
+ opts.set_xla_hlo_dump_as_html(false);
#ifdef INTEL_MKL
- flags->set_xla_cpu_use_mkl_dnn(true);
+ opts.set_xla_cpu_use_mkl_dnn(true);
#endif // INTEL_MKL
- flags->set_xla_gpu_max_kernel_unroll_factor(4);
+ opts.set_xla_gpu_max_kernel_unroll_factor(4);
// Set cudnn batchnorm off by default; it does not provide a performance win
// on average.
- flags->set_xla_gpu_use_cudnn_batchnorm(false);
+ opts.set_xla_gpu_use_cudnn_batchnorm(false);
// Run all GPU work on one stream by default. Using multiple streams
// increases memory usage and we lack strong motivating benchmarks for tuning
// the heuristics needed to decide when to run on multiple streams. See
// b/77879207.
- flags->set_xla_gpu_disable_multi_streaming(true);
+ opts.set_xla_gpu_disable_multi_streaming(true);
// TODO(jlebar): Disable fastmath once doing so is not a performance
// regression.
- flags->set_xla_cpu_enable_fast_math(true);
- flags->set_xla_gpu_enable_fast_min_max(true);
+ opts.set_xla_cpu_enable_fast_math(true);
+ opts.set_xla_gpu_enable_fast_min_max(true);
- flags->set_xla_force_host_platform_device_count(1);
+ opts.set_xla_force_host_platform_device_count(1);
+ return opts;
}
+static DebugOptions* flag_values;
+static std::vector<tensorflow::Flag>* flag_objects;
+static std::once_flag flags_init;
+
// Allocates flag_values and flag_objects; this function must not be called more
// than once - its call done via call_once.
-void AllocateFlags() {
- flag_values = new DebugOptions;
-
- SetDebugOptionsDefaults(flag_values);
+static void AllocateFlags() {
+ flag_values = new DebugOptions(DefaultDebugOptionsIgnoringFlags());
// Returns a lambda that calls "member_setter" on "flag_values" with the
// argument passed in to the lambda.
@@ -133,6 +133,11 @@
bool_setter_for(&DebugOptions::set_xla_hlo_dump_as_graphdef),
flag_values->xla_hlo_dump_as_graphdef(),
"Dump HLO graphs as TensorFlow GraphDefs."),
+ tensorflow::Flag("xla_hlo_dump_as_html",
+ bool_setter_for(&DebugOptions::set_xla_hlo_dump_as_html),
+ flag_values->xla_hlo_dump_as_html(),
+ "Dump HLO graphs as an HTML (DOT rendered into SVG "
+ "inlined in HTML)."),
tensorflow::Flag(
"xla_hlo_graph_sharding_color",
bool_setter_for(&DebugOptions::set_xla_hlo_graph_sharding_color),
@@ -203,6 +208,16 @@
"must exactly match the passes' names; no whitespace around "
"commas."),
tensorflow::Flag(
+ "xla_disable_all_hlo_passes",
+ bool_setter_for(&DebugOptions::set_xla_disable_all_hlo_passes), false,
+ "Disables all HLO passes. Notes that some passes are necessary for "
+ "correctness and the invariants that must be satisfied by 'fully "
+ "optimized' HLO are different for different devices and may change "
+ "over time. The only 'guarantee', such as it is, is that if you "
+ "compile XLA and dump the optimized HLO for some graph, you should "
+ "be able to run it again on the same device with the same build of "
+ "XLA."),
+ tensorflow::Flag(
"xla_embed_ir_in_executable",
bool_setter_for(&DebugOptions::set_xla_embed_ir_in_executable),
flag_values->xla_embed_ir_in_executable(),
@@ -334,12 +349,16 @@
"overhead from context switching but we let the user override this "
"behavior to help run tests on the host that run models in parallel "
"across multiple devices."),
+ tensorflow::Flag(
+ "xla_gpu_disable_ptxas_optimizations",
+ bool_setter_for(
+ &DebugOptions::set_xla_gpu_disable_ptxas_optimizations),
+ flag_values->xla_gpu_disable_ptxas_optimizations(),
+ "In XLA:GPU run ptxas in -O0 (default is -O3)."),
});
ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", *flag_objects);
}
-} // namespace
-
void AppendDebugOptionsFlags(std::vector<tensorflow::Flag>* flag_list) {
std::call_once(flags_init, &AllocateFlags);
flag_list->insert(flag_list->end(), flag_objects->begin(),
diff --git a/tensorflow/compiler/xla/debug_options_flags.h b/tensorflow/compiler/xla/debug_options_flags.h
index 60e59ab..dbf86a4 100644
--- a/tensorflow/compiler/xla/debug_options_flags.h
+++ b/tensorflow/compiler/xla/debug_options_flags.h
@@ -29,7 +29,10 @@
// Fetches a DebugOptions proto message from flags provided to the program.
// Flags must be registered with the flags parser using AppendDebugOptionsFlags
// first.
-xla::DebugOptions GetDebugOptionsFromFlags();
+DebugOptions GetDebugOptionsFromFlags();
+
+// Gets a DebugOptions proto that reflects the defaults as if no flags were set.
+DebugOptions DefaultDebugOptionsIgnoringFlags();
} // namespace xla
diff --git a/tensorflow/compiler/xla/g3doc/_book.yaml b/tensorflow/compiler/xla/g3doc/_book.yaml
index 12b7094..267701e 100644
--- a/tensorflow/compiler/xla/g3doc/_book.yaml
+++ b/tensorflow/compiler/xla/g3doc/_book.yaml
@@ -31,3 +31,5 @@
- title: XLA compile API
path: /xla/tutorials/xla_compile
status: experimental
+
+- include: /_upper_tabs_right.yaml
diff --git a/tensorflow/compiler/xla/g3doc/operation_semantics.md b/tensorflow/compiler/xla/g3doc/operation_semantics.md
index d888b1f..002ebc3 100644
--- a/tensorflow/compiler/xla/g3doc/operation_semantics.md
+++ b/tensorflow/compiler/xla/g3doc/operation_semantics.md
@@ -38,25 +38,25 @@
It has two phases:
1. the scatter phase. On each core, the operand is split into `split_count`
- number of blocks along the `split_dimensions`, and the blocks are scattered
- to all cores, e.g., the ith block is send to the ith core.
+number of blocks along the `split_dimensions`, and the blocks are scattered
+to all cores, e.g., the ith block is send to the ith core.
2. the gather phase. Each core concatenates the received blocks along the
- `concat_dimension`.
+`concat_dimension`.
The participating cores can be configured by:
- `replica_groups`: each ReplicaGroup contains a list of replica id. If empty,
- all replicas belong to one group in the order of 0 - (n-1). Alltoall will be
- applied within subgroups in the specified order. For example, replica
- groups = {{1,2,3},{4,5,0}} means, an Alltoall will be applied within replica
- 1, 2, 3, and in the gather phase, the received blocks will be concatenated
- in the order of 1, 2, 3; another Alltoall will be applied within replica 4,
- 5, 0, and the concatenation order is 4, 5, 0.
+all replicas belong to one group in the order of 0 - (n-1). Alltoall will be
+applied within subgroups in the specified order. For example, replica
+groups = {{1,2,3},{4,5,0}} means, an Alltoall will be applied within replica
+1, 2, 3, and in the gather phase, the received blocks will be concatenated
+in the order of 1, 2, 3; another Alltoall will be applied within replica 4,
+5, 0, and the concatenation order is 4, 5, 0.
Prerequisites:
- The dimension size of the operand on the split_dimension is divisible by
- split_count.
+split_count.
- The operand's shape is not tuple.
<b> `AllToAll(operand, split_dimension, concat_dimension, split_count,
@@ -93,7 +93,7 @@
```
<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
- <img style="width:100%" src="./images/ops_alltoall.png">
+<img style="width:100%" src="./images/ops_alltoall.png">
</div>
In this example, there are 4 cores participating the Alltoall. On each core, the
@@ -387,34 +387,34 @@
```
let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}},
- {{20, 21, 22}, {25, 26, 27}},
- {{30, 31, 32}, {35, 36, 37}},
- {{40, 41, 42}, {45, 46, 47}}};
+{{20, 21, 22}, {25, 26, 27}},
+{{30, 31, 32}, {35, 36, 37}},
+{{40, 41, 42}, {45, 46, 47}}};
// Collapse to a single dimension, leaving one dimension.
let v012 = Collapse(v, {0,1,2});
then v012 == f32[24] {10, 11, 12, 15, 16, 17,
- 20, 21, 22, 25, 26, 27,
- 30, 31, 32, 35, 36, 37,
- 40, 41, 42, 45, 46, 47};
+20, 21, 22, 25, 26, 27,
+30, 31, 32, 35, 36, 37,
+40, 41, 42, 45, 46, 47};
// Collapse the two lower dimensions, leaving two dimensions.
let v01 = Collapse(v, {0,1});
then v01 == f32[4x6] {{10, 11, 12, 15, 16, 17},
- {20, 21, 22, 25, 26, 27},
- {30, 31, 32, 35, 36, 37},
- {40, 41, 42, 45, 46, 47}};
+{20, 21, 22, 25, 26, 27},
+{30, 31, 32, 35, 36, 37},
+{40, 41, 42, 45, 46, 47}};
// Collapse the two higher dimensions, leaving two dimensions.
let v12 = Collapse(v, {1,2});
then v12 == f32[8x3] {{10, 11, 12},
- {15, 16, 17},
- {20, 21, 22},
- {25, 26, 27},
- {30, 31, 32},
- {35, 36, 37},
- {40, 41, 42},
- {45, 46, 47}};
+{15, 16, 17},
+{20, 21, 22},
+{25, 26, 27},
+{30, 31, 32},
+{35, 36, 37},
+{40, 41, 42},
+{45, 46, 47}};
```
@@ -441,9 +441,9 @@
Note that there are the following restrictions on the `source_target_pair`:
- Any two pairs should not have the same target replica id, and they should
- not have the same source replica id.
+not have the same source replica id.
- If a replica id is not a target in any pair, then the output on that replica
- is a tensor consists of 0(s) with the same shape as the input.
+is a tensor consists of 0(s) with the same shape as the input.
## Concatenate
@@ -480,25 +480,25 @@
```
let a = {
- {1, 2},
- {3, 4},
- {5, 6},
+{1, 2},
+{3, 4},
+{5, 6},
};
let b = {
- {7, 8},
+{7, 8},
};
Concat({a, b}, 0)
>>> {
- {1, 2},
- {3, 4},
- {5, 6},
- {7, 8},
+{1, 2},
+{3, 4},
+{5, 6},
+{7, 8},
}
```
Diagram:
<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
- <img style="width:100%" src="./images/ops_concatenate.png">
+<img style="width:100%" src="./images/ops_concatenate.png">
</div>
## Conditional
@@ -566,20 +566,20 @@
The n+2 dimensions are, in this order:
* `batch`: Each coordinate in this dimension represents an independent input
- for which convolution is carried out.
+for which convolution is carried out.
* `z/depth/features`: Each (y,x) position in the base area has a vector
- associated to it, which goes into this dimension.
+associated to it, which goes into this dimension.
* `spatial_dims`: Describes the `n` spatial dimensions that define the base
- area that the window moves across.
+area that the window moves across.
The `rhs` argument is a rank n+2 array describing the convolutional
filter/kernel/window. The dimensions are, in this order:
* `output-z`: The `z` dimension of the output.
* `input-z`: The size of this dimension times `feature_group_count` should
- equal the size of the `z` dimension in lhs.
+equal the size of the `z` dimension in lhs.
* `spatial_dims`: Describes the `n` spatial dimensions that define the n-d
- window that moves across the base area.
+window that moves across the base area.
The `window_strides` argument specifies the stride of the convolutional window
in the spatial dimensions. For example, if the stride in the first spatial
@@ -633,7 +633,7 @@
* `batch`: Same size as `batch` on the input (`lhs`).
* `z`: Same size as `output-z` on the kernel (`rhs`).
* `spatial_dims`: One value for each valid placement of the convolutional
- window.
+window.
The valid placements of the convolutional window are determined by the strides
and the size of the base area after padding.
@@ -658,15 +658,15 @@
```
for (b, oz, oy, ox) { // output coordinates
- value = 0;
- for (iz, ky, kx) { // kernel coordinates and input z
- iy = oy*stride_y + ky - pad_low_y;
- ix = ox*stride_x + kx - pad_low_x;
- if ((iy, ix) inside the base area considered without padding) {
- value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
- }
- }
- output(b, oz, oy, ox) = value;
+value = 0;
+for (iz, ky, kx) { // kernel coordinates and input z
+iy = oy*stride_y + ky - pad_low_y;
+ix = ox*stride_x + kx - pad_low_x;
+if ((iy, ix) inside the base area considered without padding) {
+value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
+}
+}
+output(b, oz, oy, ox) = value;
}
```
@@ -777,19 +777,19 @@
```
extern "C" void myfunc(void* out, void** in) {
- float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
- float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
- EXPECT_EQ(1, x[0]);
- EXPECT_EQ(2, x[1]);
- EXPECT_EQ(10, y[0][0]);
- EXPECT_EQ(20, y[0][1]);
- EXPECT_EQ(30, y[0][2]);
- EXPECT_EQ(40, y[1][0]);
- EXPECT_EQ(50, y[1][1]);
- EXPECT_EQ(60, y[1][2]);
- float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
- z[0][0] = x[1] + y[1][0];
- // ...
+float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
+float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
+EXPECT_EQ(1, x[0]);
+EXPECT_EQ(2, x[1]);
+EXPECT_EQ(10, y[0][0]);
+EXPECT_EQ(20, y[0][1]);
+EXPECT_EQ(30, y[0][2]);
+EXPECT_EQ(40, y[1][0]);
+EXPECT_EQ(50, y[1][1]);
+EXPECT_EQ(60, y[1][2]);
+float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
+z[0][0] = x[1] + y[1][0];
+// ...
}
```
@@ -864,17 +864,17 @@
```
lhs = { {1.0, 2.0, 3.0},
- {4.0, 5.0, 6.0} }
+{4.0, 5.0, 6.0} }
rhs = { {1.0, 1.0, 1.0},
- {2.0, 2.0, 2.0} }
+{2.0, 2.0, 2.0} }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);
DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
- {15.0, 30.0} }
+{15.0, 30.0} }
```
Associated batch dimension numbers from the 'lhs' and 'rhs' must have the same
@@ -886,14 +886,14 @@
```
lhs = { { {1.0, 2.0},
- {3.0, 4.0} },
- { {5.0, 6.0},
- {7.0, 8.0} } }
+{3.0, 4.0} },
+{ {5.0, 6.0},
+{7.0, 8.0} } }
rhs = { { {1.0, 0.0},
- {0.0, 1.0} },
- { {1.0, 0.0},
- {0.0, 1.0} } }
+{0.0, 1.0} },
+{ {1.0, 0.0},
+{0.0, 1.0} } }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
@@ -902,9 +902,9 @@
dnums.add_rhs_batch_dimensions(0);
DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
- {3.0, 4.0} },
- { {5.0, 6.0},
- {7.0, 8.0} } }
+{3.0, 4.0} },
+{ {5.0, 6.0},
+{7.0, 8.0} } }
```
| Input | Output | Semantics |
@@ -963,22 +963,22 @@
let s = {2}
DynamicSlice(a, s, {2}) produces:
- {2.0, 3.0}
+{2.0, 3.0}
```
2-dimensional example:
```
let b =
- { {0.0, 1.0, 2.0},
- {3.0, 4.0, 5.0},
- {6.0, 7.0, 8.0},
- {9.0, 10.0, 11.0} }
+{ {0.0, 1.0, 2.0},
+{3.0, 4.0, 5.0},
+{6.0, 7.0, 8.0},
+{9.0, 10.0, 11.0} }
let s = {2, 1}
DynamicSlice(b, s, {2, 2}) produces:
- { { 7.0, 8.0},
- {10.0, 11.0} }
+{ { 7.0, 8.0},
+{10.0, 11.0} }
```
## DynamicUpdateSlice
@@ -1027,29 +1027,29 @@
let s = {2}
DynamicUpdateSlice(a, u, s) produces:
- {0.0, 1.0, 5.0, 6.0, 4.0}
+{0.0, 1.0, 5.0, 6.0, 4.0}
```
2-dimensional example:
```
let b =
- { {0.0, 1.0, 2.0},
- {3.0, 4.0, 5.0},
- {6.0, 7.0, 8.0},
- {9.0, 10.0, 11.0} }
+{ {0.0, 1.0, 2.0},
+{3.0, 4.0, 5.0},
+{6.0, 7.0, 8.0},
+{9.0, 10.0, 11.0} }
let u =
- { {12.0, 13.0},
- {14.0, 15.0},
- {16.0, 17.0} }
+{ {12.0, 13.0},
+{14.0, 15.0},
+{16.0, 17.0} }
let s = {1, 1}
DynamicUpdateSlice(b, u, s) produces:
- { {0.0, 1.0, 2.0},
- {3.0, 12.0, 13.0},
- {6.0, 14.0, 15.0},
- {9.0, 16.0, 17.0} }
+{ {0.0, 1.0, 2.0},
+{3.0, 12.0, 13.0},
+{6.0, 14.0, 15.0},
+{9.0, 16.0, 17.0} }
```
## Element-wise binary arithmetic operations
@@ -1235,42 +1235,42 @@
The bounds for the output array along dimension `i` is computed as follows:
- 1. If `i` is present in `batch_dims` (i.e. is equal to `batch_dims[k]` for
- some `k`) then we pick the corresponding dimension bounds out of
- `start_indices.shape`, skipping `index_vector_dim` (i.e. pick
- `start_indices.shape.dims`[`k`] if `k` < `index_vector_dim` and
- `start_indices.shape.dims`[`k`+`1`] otherwise).
+1. If `i` is present in `batch_dims` (i.e. is equal to `batch_dims[k]` for
+some `k`) then we pick the corresponding dimension bounds out of
+`start_indices.shape`, skipping `index_vector_dim` (i.e. pick
+`start_indices.shape.dims`[`k`] if `k` < `index_vector_dim` and
+`start_indices.shape.dims`[`k`+`1`] otherwise).
- 2. If `i` is present in `offset_dims` (i.e. equal to `offset_dims`[`k`] for
- some `k`) then we pick the corresponding bound out of `slice_sizes` after
- accounting for `collapsed_slice_dims` (i.e. we pick
- `adjusted_slice_sizes`[`k`] where `adjusted_slice_sizes` is `slice_sizes`
- with the bounds at indices `collapsed_slice_dims` removed).
+2. If `i` is present in `offset_dims` (i.e. equal to `offset_dims`[`k`] for
+some `k`) then we pick the corresponding bound out of `slice_sizes` after
+accounting for `collapsed_slice_dims` (i.e. we pick
+`adjusted_slice_sizes`[`k`] where `adjusted_slice_sizes` is `slice_sizes`
+with the bounds at indices `collapsed_slice_dims` removed).
Formally, the operand index `In` corresponding to an output index `Out` is
computed as follows:
- 1. Let `G` = { `Out`[`k`] for `k` in `batch_dims` }. Use `G` to slice out
- vector `S` such that `S`[`i`] = `start_indices`[Combine(`G`, `i`)] where
- Combine(A, b) inserts b at position `index_vector_dim` into A. Note that
- this is well defined even if `G` is empty -- if `G` is empty then `S` =
- `start_indices`.
+1. Let `G` = { `Out`[`k`] for `k` in `batch_dims` }. Use `G` to slice out
+vector `S` such that `S`[`i`] = `start_indices`[Combine(`G`, `i`)] where
+Combine(A, b) inserts b at position `index_vector_dim` into A. Note that
+this is well defined even if `G` is empty -- if `G` is empty then `S` =
+`start_indices`.
- 2. Create a starting index, `S`<sub>`in`</sub>, into `operand` using `S` by
- scattering `S` using `start_index_map`. More precisely:
- 1. `S`<sub>`in`</sub>[`start_index_map`[`k`]] = `S`[`k`] if `k` <
- `start_index_map.size`.
- 2. `S`<sub>`in`</sub>[`_`] = `0` otherwise.
+2. Create a starting index, `S`<sub>`in`</sub>, into `operand` using `S` by
+scattering `S` using `start_index_map`. More precisely:
+1. `S`<sub>`in`</sub>[`start_index_map`[`k`]] = `S`[`k`] if `k` <
+`start_index_map.size`.
+2. `S`<sub>`in`</sub>[`_`] = `0` otherwise.
- 3. Create an index `O`<sub>`in`</sub> into `operand` by scattering the indices
- at the offset dimensions in `Out` according to the `collapsed_slice_dims`
- set. More precisely:
- 1. `O`<sub>`in`</sub>[`expand_offset_dims`(`k`)] =
- `Out`[`offset_dims`[`k`]] if `k` < `offset_dims.size`
- (`expand_offset_dims` is defined below).
- 2. `O`<sub>`in`</sub>[`_`] = `0` otherwise.
- 4. `In` is `O`<sub>`in`</sub> + `S`<sub>`in`</sub> where + is element-wise
- addition.
+3. Create an index `O`<sub>`in`</sub> into `operand` by scattering the indices
+at the offset dimensions in `Out` according to the `collapsed_slice_dims`
+set. More precisely:
+1. `O`<sub>`in`</sub>[`expand_offset_dims`(`k`)] =
+`Out`[`offset_dims`[`k`]] if `k` < `offset_dims.size`
+(`expand_offset_dims` is defined below).
+2. `O`<sub>`in`</sub>[`_`] = `0` otherwise.
+4. `In` is `O`<sub>`in`</sub> + `S`<sub>`in`</sub> where + is element-wise
+addition.
`expand_offset_dims` is the monotonic function with domain [`0`, `offset.size`)
and range [`0`, `operand.rank`) \ `collapsed_slice_dims`. So if, e.g.,
@@ -1282,21 +1282,21 @@
Informally, every index `Out` in the output array corresponds to an element `E`
in the operand array, computed as follows:
- - We use the batch dimensions in `Out` to look up a starting index from
- `start_indices`.
+- We use the batch dimensions in `Out` to look up a starting index from
+`start_indices`.
- - We use `start_index_map` to map the starting index (which may have size less
- than operand.rank) to a "full" starting index into operand.
+- We use `start_index_map` to map the starting index (which may have size less
+than operand.rank) to a "full" starting index into operand.
- - We dynamic-slice out a slice with size `slice_sizes` using the full starting
- index.
+- We dynamic-slice out a slice with size `slice_sizes` using the full starting
+index.
- - We reshape the slice by collapsing the `collapsed_slice_dims` dimensions.
- Since all collapsed slice dimensions have to have bound 1 this reshape is
- always legal.
+- We reshape the slice by collapsing the `collapsed_slice_dims` dimensions.
+Since all collapsed slice dimensions have to have bound 1 this reshape is
+always legal.
- - We use the offset dimensions in `Out` to index into this slice to get the
- input element, `E`, corresponding to output index `Out`.
+- We use the offset dimensions in `Out` to index into this slice to get the
+input element, `E`, corresponding to output index `Out`.
`index_vector_dim` is set to `start_indices.rank` - `1` in all of the
examples that follow. More interesting values for `index_vector_dim` does not
@@ -1315,7 +1315,7 @@
way:
<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
- <img style="width:100%" src="./images/ops_xla_gather_0.svg">
+<img style="width:100%" src="./images/ops_xla_gather_0.svg">
</div>
We first select an (`X`,`Y`) vector from the gather indices array using `G`.
@@ -1334,7 +1334,7 @@
would translate indices like this:
<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
- <img style="width:100%" src="./images/ops_xla_gather_1.svg">
+<img style="width:100%" src="./images/ops_xla_gather_1.svg">
</div>
Again, this acts as a batch dynamic slice `G`<sub>`0`</sub> and
@@ -1343,27 +1343,27 @@
The gather operation in XLA generalizes the informal semantics outlined above in
the following ways:
- 1. We can configure which dimensions in the output shape are the offset
- dimensions (dimensions containing `O`<sub>`0`</sub>, `O`<sub>`1`</sub> in
- the last example). The output batch dimensions (dimensions containing
- `G`<sub>`0`</sub>, `G`<sub>`1`</sub> in the last example) are defined to be
- the output dimensions that are not offset dimensions.
+1. We can configure which dimensions in the output shape are the offset
+dimensions (dimensions containing `O`<sub>`0`</sub>, `O`<sub>`1`</sub> in
+the last example). The output batch dimensions (dimensions containing
+`G`<sub>`0`</sub>, `G`<sub>`1`</sub> in the last example) are defined to be
+the output dimensions that are not offset dimensions.
- 2. The number of output offset dimensions explicitly present in the output
- shape may be smaller than the input rank. These "missing" dimensions, which
- are listed explicitly as `collapsed_slice_dims`, must have a slice size of
- `1`. Since they have a slice size of `1` the only valid index for them is
- `0` and eliding them does not introduce ambiguity.
+2. The number of output offset dimensions explicitly present in the output
+shape may be smaller than the input rank. These "missing" dimensions, which
+are listed explicitly as `collapsed_slice_dims`, must have a slice size of
+`1`. Since they have a slice size of `1` the only valid index for them is
+`0` and eliding them does not introduce ambiguity.
- 3. The slice extracted from the "Gather Indices" array ((`X`, `Y`) in the last
- example) may have fewer elements than the input array rank, and an explicit
- mapping dictates how the index should be expanded to have the same rank as
- the input.
+3. The slice extracted from the "Gather Indices" array ((`X`, `Y`) in the last
+example) may have fewer elements than the input array rank, and an explicit
+mapping dictates how the index should be expanded to have the same rank as
+the input.
As a final example, we use (2) and (3) to implement `tf.gather_nd`:
<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
- <img style="width:100%" src="./images/ops_xla_gather_2.svg">
+<img style="width:100%" src="./images/ops_xla_gather_2.svg">
</div>
`G`<sub>`0`</sub> and `G`<sub>`1`</sub> are used to slice out a starting index
@@ -1442,11 +1442,11 @@
```
result1 = while (condition, init = init_value) {
- Infeed(shape)
+Infeed(shape)
}
result2 = while (condition, init = result1) {
- Infeed(shape)
+Infeed(shape)
}
```
@@ -1464,7 +1464,9 @@
Builds a constant literal on device rather than a potentially large host
transfer. Creates a rank 1 array of values starting at zero and incrementing by
-one.
+one. For floating-point types, the produced array is equivalent to
+`ConvertElementType(Iota(...))` where the `Iota` is of integral type and the
+conversion is to the floating-point type.
Arguments | Type | Semantics
---------------- | --------------- | ------------------------------------
diff --git a/tensorflow/compiler/xla/layout.cc b/tensorflow/compiler/xla/layout.cc
new file mode 100644
index 0000000..e3b5fcd
--- /dev/null
+++ b/tensorflow/compiler/xla/layout.cc
@@ -0,0 +1,96 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/layout.h"
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/compiler/xla/layout_util.h"
+
+namespace xla {
+
+TileProto Tile::ToProto() const {
+ TileProto tile_proto;
+ for (int64 i : dimensions()) {
+ tile_proto.add_dimensions(i);
+ }
+ return tile_proto;
+}
+
+string Tile::ToString() const {
+ return absl::StrCat("(", absl::StrJoin(dimensions(), ","), ")");
+}
+
+/* static */ Layout Layout::CreateFromProto(const LayoutProto& proto) {
+ Layout layout;
+ layout.set_format(proto.format());
+ layout.minor_to_major_.reserve(proto.minor_to_major_size());
+ for (const int64 dimension : proto.minor_to_major()) {
+ layout.add_minor_to_major(dimension);
+ }
+ layout.set_max_sparse_elements(proto.max_sparse_elements());
+ for (const TileProto& tile_proto : proto.tiles()) {
+ *layout.add_tiles() = Tile::CreateFromProto(tile_proto);
+ }
+ layout.set_element_size_in_bits(proto.element_size_in_bits());
+ return layout;
+}
+
+LayoutProto Layout::ToProto() const {
+ LayoutProto proto;
+ proto.set_format(format_);
+ proto.mutable_minor_to_major()->Reserve(minor_to_major_size());
+ for (const int64 dimension : minor_to_major()) {
+ proto.add_minor_to_major(dimension);
+ }
+ proto.set_max_sparse_elements(max_sparse_elements_);
+ for (const Tile& tile : tiles()) {
+ *proto.add_tiles() = tile.ToProto();
+ }
+ proto.set_element_size_in_bits(element_size_in_bits());
+ return proto;
+}
+
+string Layout::ToString() const {
+ // TODO(b/119839262): Emit tiles in string.
+ if (format() == SPARSE) {
+ return absl::StrCat("sparse{", max_sparse_elements(), "}");
+ } else if (format() == DENSE) {
+ return absl::StrCat("{", absl::StrJoin(minor_to_major(), ","), "}");
+ } else {
+ CHECK_EQ(format(), INVALID_FORMAT);
+ return "invalid{}";
+ }
+}
+
+bool Layout::operator==(const Layout& other) const {
+ return (other.format() == format() &&
+ other.minor_to_major() == minor_to_major() &&
+ other.element_size_in_bits() == element_size_in_bits() &&
+ other.max_sparse_elements() == max_sparse_elements() &&
+ other.tiles() == tiles());
+}
+
+std::ostream& operator<<(std::ostream& out, const Tile& tile) {
+ out << tile.ToString();
+ return out;
+}
+
+std::ostream& operator<<(std::ostream& out, const Layout& layout) {
+ out << layout.ToString();
+ return out;
+}
+
+} // namespace xla
diff --git a/tensorflow/compiler/xla/layout.h b/tensorflow/compiler/xla/layout.h
new file mode 100644
index 0000000..313368c
--- /dev/null
+++ b/tensorflow/compiler/xla/layout.h
@@ -0,0 +1,187 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_LAYOUT_H_
+#define TENSORFLOW_COMPILER_XLA_LAYOUT_H_
+
+#include <vector>
+
+#include "absl/types/span.h"
+
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/compiler/xla/util.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace xla {
+
+// Describes a tile used in tiling-based layout. Refer to
+// g3doc/third_party/tensorflow/compiler/xla/g3doc/layout_with_tiling.md for
+// details.
+class Tile {
+ public:
+ Tile() = default;
+ explicit Tile(absl::Span<const int64> dimensions)
+ : dimensions_(dimensions.begin(), dimensions.end()) {}
+
+ // De/Serialize a Tile to and from a TileProto.
+ static Tile CreateFromProto(const TileProto& tile_proto) {
+ return Tile(AsInt64Slice(tile_proto.dimensions()));
+ }
+ TileProto ToProto() const;
+
+ bool operator==(const Tile& other) const {
+ return dimensions() == other.dimensions();
+ }
+ bool operator!=(const Tile& other) const { return !(*this == other); }
+
+ string ToString() const;
+
+ // Returns the bound of the tile in the given dimension index.
+ int64 dimension(int i) const { return dimensions_.at(i); }
+
+ // Returns the dimensions of the tile.
+ const std::vector<int64>& dimensions() const { return dimensions_; }
+
+ private:
+ // The bounds of the tile.
+ std::vector<int64> dimensions_;
+};
+
+class Layout {
+ public:
+ Layout() = default;
+
+ // Constructs a dense layout with the given minor-to-major order.
+ explicit Layout(absl::Span<const int64> minor_to_major)
+ : format_(DENSE),
+ minor_to_major_(minor_to_major.begin(), minor_to_major.end()) {}
+
+ // Constructs a dense tiled layout with the given minor-to-major order and
+ // tiles.
+ Layout(absl::Span<const int64> minor_to_major, absl::Span<const Tile> tiles)
+ : format_(DENSE),
+ minor_to_major_(minor_to_major.begin(), minor_to_major.end()),
+ tiles_(tiles.begin(), tiles.end()) {}
+
+ // Construct a shape from a LayoutProto.
+ static Layout CreateFromProto(const LayoutProto& proto);
+
+ // Returns a LayoutProto representation of the Layout.
+ LayoutProto ToProto() const;
+
+ // Returns a human-readable string that represents this layout.
+ string ToString() const;
+
+ bool operator==(const Layout& other) const;
+ bool operator!=(const Layout& other) const { return !(*this == other); }
+
+ // The following methods mirror the protobuf generated code interface for the
+ // message LayoutProto. This enabled easy migration of this data structure
+ // from a proto to a proper C++ class.
+ //
+ // TODO(b/29771030): Replace or augment these methods with a more ergonomic
+ // interface.
+
+ // Methods for accessing the format.
+ Format format() const { return format_; }
+ Layout& set_format(Format value) {
+ format_ = value;
+ return *this;
+ }
+
+ // Methods for accessing the minor-to-major array.
+ int minor_to_major_size() const { return minor_to_major_.size(); }
+ int64 minor_to_major(int index) const { return minor_to_major_.at(index); }
+ Layout& set_minor_to_major(int index, int64 value) {
+ minor_to_major_.at(index) = value;
+ return *this;
+ }
+ Layout& add_minor_to_major(int64 value) {
+ minor_to_major_.push_back(value);
+ return *this;
+ }
+ Layout& clear_minor_to_major() {
+ minor_to_major_.clear();
+ return *this;
+ }
+ const std::vector<int64>& minor_to_major() const { return minor_to_major_; }
+ std::vector<int64>* mutable_minor_to_major() { return &minor_to_major_; }
+
+ // Methods for accessing the tile field.
+ int tiles_size() const { return tiles_.size(); }
+ const Tile& tiles(int index) const { return tiles_.at(index); }
+ Tile* mutable_tiles(int index) { return &tiles_.at(index); }
+ Tile* add_tiles() {
+ tiles_.push_back(Tile());
+ return &tiles_.back();
+ }
+ Layout& clear_tiles() {
+ tiles_.clear();
+ return *this;
+ }
+ const std::vector<Tile>& tiles() const { return tiles_; }
+ std::vector<Tile>* mutable_tiles() { return &tiles_; }
+
+ // Methods for accessing the int64 fields.
+ int64 max_sparse_elements() const { return max_sparse_elements_; }
+ Layout& set_max_sparse_elements(int64 value) {
+ max_sparse_elements_ = value;
+ return *this;
+ }
+ int64 element_size_in_bits() const { return element_size_in_bits_; }
+ Layout& set_element_size_in_bits(int64 value) {
+ element_size_in_bits_ = value;
+ return *this;
+ }
+
+ void Swap(Layout* other) {
+ using std::swap;
+ swap(*this, *other);
+ }
+
+ void Clear() {
+ format_ = INVALID_FORMAT;
+ minor_to_major_.clear();
+ max_sparse_elements_ = 0;
+ element_size_in_bits_ = 0;
+ }
+
+ public:
+ // The format of this layout.
+ Format format_ = INVALID_FORMAT;
+
+ // Sequence of dimension numbers, from minor (fastest varying index) to major
+ // (slowest varying index).
+ std::vector<int64> minor_to_major_;
+
+ // The maximum number of elements that can be stored for SPARSE formats. This
+ // can be used to determine the maximum size in bytes of arrays stored in
+ // memory. This field must be zero unless the format is SPARSE.
+ int64 max_sparse_elements_ = 0;
+
+ // The number of bits used to store an individual array element.
+ int64 element_size_in_bits_ = 0;
+
+ // The tiles used in tiling-based layout.
+ std::vector<Tile> tiles_;
+};
+
+std::ostream& operator<<(std::ostream& out, const Tile& Tile);
+std::ostream& operator<<(std::ostream& out, const Layout& layout);
+
+} // namespace xla
+
+#endif // TENSORFLOW_COMPILER_XLA_LAYOUT_H_
diff --git a/tensorflow/compiler/xla/layout_test.cc b/tensorflow/compiler/xla/layout_test.cc
new file mode 100644
index 0000000..fb6abd3
--- /dev/null
+++ b/tensorflow/compiler/xla/layout_test.cc
@@ -0,0 +1,104 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/layout.h"
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/compiler/xla/layout_util.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/status_macros.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/test_helpers.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/compiler/xla/util.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+
+namespace xla {
+namespace {
+
+class LayoutTest : public ::testing::Test {};
+
+TEST_F(LayoutTest, ToString) {
+ EXPECT_EQ(Layout().ToString(), "invalid{}");
+ EXPECT_EQ(Layout({4, 5, 6}).ToString(), "{4,5,6}");
+ EXPECT_EQ(Layout().set_format(SPARSE).set_max_sparse_elements(123).ToString(),
+ "sparse{123}");
+ EXPECT_EQ(Layout({4, 5, 6}).ToString(), "{4,5,6}");
+ EXPECT_EQ(Layout({3, 2, 1, 0}, {Tile({42, 123}), Tile({4, 5})}).ToString(),
+ "{3,2,1,0}");
+ EXPECT_EQ(
+ Layout({1, 0}, {Tile({2, 55})}).set_element_size_in_bits(42).ToString(),
+ "{1,0}");
+}
+
+TEST_F(LayoutTest, StreamOut) {
+ {
+ std::ostringstream oss;
+ oss << Tile({7, 8});
+ EXPECT_EQ(oss.str(), "(7,8)");
+ }
+
+ {
+ std::ostringstream oss;
+ oss << Layout({0, 1, 2});
+ EXPECT_EQ(oss.str(), "{0,1,2}");
+ }
+}
+
+TEST_F(LayoutTest, SparseLayoutMaxElements) {
+ EXPECT_EQ(LayoutUtil::MaxSparseElements(LayoutUtil::MakeSparseLayout(101)),
+ 101);
+}
+
+TEST_F(LayoutTest, Equality) {
+ EXPECT_EQ(Layout(), Layout());
+ const std::vector<int64> empty_dims;
+ EXPECT_EQ(Layout(empty_dims), Layout(empty_dims));
+ EXPECT_NE(Layout(), Layout(empty_dims));
+ EXPECT_EQ(Layout({0, 1, 2, 3}), Layout({0, 1, 2, 3}));
+ EXPECT_NE(Layout({0, 1, 2, 3}), Layout({0, 1, 2}));
+ EXPECT_EQ(Layout({0, 1, 2}, {Tile({42, 44})}),
+ Layout({0, 1, 2}, {Tile({42, 44})}));
+ EXPECT_NE(Layout({0, 1, 2}, {Tile({42, 44})}),
+ Layout({0, 1, 2}, {Tile({42, 45})}));
+ EXPECT_NE(Layout({0, 1, 2}, {Tile({42, 44})}), Layout({0, 1, 2, 3}));
+ EXPECT_EQ(Layout({0, 1, 2}).set_element_size_in_bits(33),
+ Layout({0, 1, 2}).set_element_size_in_bits(33));
+ EXPECT_NE(Layout({0, 1, 2}).set_element_size_in_bits(33),
+ Layout({0, 1, 2}).set_element_size_in_bits(7));
+ EXPECT_EQ(Layout().set_format(SPARSE), Layout().set_format(SPARSE));
+ EXPECT_EQ(Layout().set_format(SPARSE).set_max_sparse_elements(42),
+ Layout().set_format(SPARSE).set_max_sparse_elements(42));
+ EXPECT_NE(Layout().set_format(SPARSE).set_max_sparse_elements(42),
+ Layout().set_format(SPARSE).set_max_sparse_elements(24));
+}
+
+TEST_F(LayoutTest, LayoutToFromProto) {
+ // Round-trips a Layout through proto de/serialization.
+ auto expect_unchanged = [](const Layout& layout) {
+ EXPECT_EQ(layout, Layout::CreateFromProto(layout.ToProto()));
+ };
+
+ expect_unchanged(Layout());
+ expect_unchanged(Layout({1, 3, 2, 0}));
+ expect_unchanged(Layout().set_format(SPARSE));
+ expect_unchanged(Layout().set_format(SPARSE).set_max_sparse_elements(123));
+ expect_unchanged(Layout({0, 1}).set_element_size_in_bits(42));
+ expect_unchanged(Layout({3, 2, 1, 0}, {Tile({42, 123}), Tile({4, 5})}));
+}
+
+} // namespace
+} // namespace xla
diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc
index dbb8138..ddccd8c 100644
--- a/tensorflow/compiler/xla/layout_util.cc
+++ b/tensorflow/compiler/xla/layout_util.cc
@@ -41,15 +41,13 @@
// Internal helper for GetDefaultLayoutForShape and SetToDefaultLayout. Sets
// minor_to_major to the value that represents the default layout.
-void SetDefaultLayoutToContainer(
- tensorflow::protobuf::RepeatedField<tensorflow::protobuf_int64>*
- minor_to_major) {
+void SetDefaultLayoutToContainer(std::vector<int64>* minor_to_major) {
// The default XLA layout is major-to-minor (dim 0 is major).
// For more information on XLA layouts, see:
// https://www.tensorflow.org/performance/xla/shapes
const int64 size = minor_to_major->size();
for (int64 i = 0; i < size; ++i) {
- minor_to_major->Set(i, size - 1 - i);
+ (*minor_to_major)[i] = size - 1 - i;
}
}
@@ -94,9 +92,8 @@
Layout CreateDefaultLayoutForRank(int64 rank) {
Layout layout;
layout.set_format(DENSE);
- tensorflow::protobuf::RepeatedField<tensorflow::protobuf_int64>*
- minor_to_major = layout.mutable_minor_to_major();
- minor_to_major->Resize(rank, 0);
+ std::vector<int64>* minor_to_major = layout.mutable_minor_to_major();
+ minor_to_major->resize(rank, 0);
SetDefaultLayoutToContainer(minor_to_major);
return layout;
}
@@ -139,9 +136,8 @@
shape->clear_layout();
} else if (ShapeUtil::IsArray(*shape)) {
shape->mutable_layout()->set_format(DENSE);
- tensorflow::protobuf::RepeatedField<tensorflow::protobuf_int64>*
- minor_to_major = shape->mutable_layout()->mutable_minor_to_major();
- minor_to_major->Resize(shape->dimensions_size(), 0);
+ auto* minor_to_major = shape->mutable_layout()->mutable_minor_to_major();
+ minor_to_major->resize(shape->dimensions_size(), 0);
SetDefaultLayoutToContainer(minor_to_major);
} else {
// Opaque, token types etc. have no layout.
@@ -210,9 +206,8 @@
}
if (layout.format() == INVALID_FORMAT || !Format_IsValid(layout.format())) {
- return InvalidArgument(
- "Layout has an invalid format (%d) in layout {%s}, shape {%s}",
- layout.format(), layout.ShortDebugString(), shape.ShortDebugString());
+ return InvalidArgument("Layout has an invalid format (%d)",
+ layout.format());
}
if (layout.format() == DENSE) {
@@ -316,7 +311,7 @@
}
/* static */ bool LayoutUtil::Equal(const Layout& lhs, const Layout& rhs) {
- return protobuf_util::ProtobufEquals(lhs, rhs);
+ return lhs == rhs;
}
/* static */ absl::Span<const int64> LayoutUtil::MinorToMajor(
@@ -358,11 +353,7 @@
}
/* static */ string LayoutUtil::HumanString(const Layout& layout) {
- if (IsSparse(layout)) {
- return absl::StrCat("sparse{", layout.max_sparse_elements(), "}");
- }
- CHECK(IsDense(layout));
- return absl::StrCat("{", absl::StrJoin(layout.minor_to_major(), ","), "}");
+ return layout.ToString();
}
namespace {
@@ -444,11 +435,6 @@
return true;
}
-std::ostream& operator<<(std::ostream& out, const Layout& layout) {
- out << LayoutUtil::HumanString(layout);
- return out;
-}
-
/*static*/ size_t LayoutUtil::Hash(const Layout& layout) {
using tensorflow::hash;
using tensorflow::Hash64Combine;
diff --git a/tensorflow/compiler/xla/layout_util.h b/tensorflow/compiler/xla/layout_util.h
index 6c298e5..609dba6 100644
--- a/tensorflow/compiler/xla/layout_util.h
+++ b/tensorflow/compiler/xla/layout_util.h
@@ -21,6 +21,7 @@
#include <string>
#include "absl/types/span.h"
+#include "tensorflow/compiler/xla/layout.h"
#include "tensorflow/compiler/xla/shape.h"
#include "tensorflow/compiler/xla/status.h"
#include "tensorflow/compiler/xla/types.h"
@@ -195,8 +196,6 @@
TF_DISALLOW_COPY_AND_ASSIGN(LayoutUtil);
};
-std::ostream& operator<<(std::ostream& out, const Layout& layout);
-
} // namespace xla
#endif // TENSORFLOW_COMPILER_XLA_LAYOUT_UTIL_H_
diff --git a/tensorflow/compiler/xla/layout_util_test.cc b/tensorflow/compiler/xla/layout_util_test.cc
index 12ce2d2..4cc94c27 100644
--- a/tensorflow/compiler/xla/layout_util_test.cc
+++ b/tensorflow/compiler/xla/layout_util_test.cc
@@ -317,17 +317,6 @@
ShapeUtil::MakeShape(F32, {10, 20, 30, 15, 25}))));
}
-TEST_F(LayoutUtilTest, SparseLayoutMaxElements) {
- EXPECT_EQ(LayoutUtil::MaxSparseElements(LayoutUtil::MakeSparseLayout(101)),
- 101);
-}
-
-TEST_F(LayoutUtilTest, StreamOut) {
- std::ostringstream oss;
- oss << LayoutUtil::MakeLayout({0, 1, 2});
- EXPECT_EQ(oss.str(), "{0,1,2}");
-}
-
TEST_F(LayoutUtilTest, ValidateLayout_ValidArrayLayout) {
Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {0, 1});
auto status =
diff --git a/tensorflow/compiler/xla/packed_literal_reader.cc b/tensorflow/compiler/xla/packed_literal_reader.cc
index 0f86f9f..339660c 100644
--- a/tensorflow/compiler/xla/packed_literal_reader.cc
+++ b/tensorflow/compiler/xla/packed_literal_reader.cc
@@ -42,8 +42,7 @@
StatusOr<Literal> PackedLiteralReader::Read(const Shape& shape,
const Layout* layout) {
VLOG(3) << "reading shape from file: " << ShapeUtil::HumanString(shape)
- << " layout: "
- << (layout == nullptr ? "<none>" : layout->ShortDebugString());
+ << " layout: " << (layout == nullptr ? "<none>" : layout->ToString());
Shape literal_shape = shape;
if (layout != nullptr) {
TF_RETURN_IF_ERROR(
diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD
index 63ac1c6..4a57b10 100644
--- a/tensorflow/compiler/xla/python/BUILD
+++ b/tensorflow/compiler/xla/python/BUILD
@@ -66,7 +66,10 @@
"//tensorflow/compiler/xla/client:local_client",
"//tensorflow/compiler/xla/client:xla_builder",
"//tensorflow/compiler/xla/client:xla_computation",
+ "//tensorflow/compiler/xla/client/lib:cholesky",
"//tensorflow/compiler/xla/client/lib:math",
+ "//tensorflow/compiler/xla/client/lib:qr",
+ "//tensorflow/compiler/xla/client/lib:triangular_solve",
"//tensorflow/compiler/xla/service:platform_util",
"//tensorflow/compiler/xla/service:shaped_buffer",
"//tensorflow/compiler/xrt:xrt_proto",
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc
index 6e2ee86..5d191f5 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.cc
+++ b/tensorflow/compiler/xla/python/local_computation_builder.cc
@@ -24,7 +24,10 @@
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/compiler/xla/client/lib/cholesky.h"
#include "tensorflow/compiler/xla/client/lib/math.h"
+#include "tensorflow/compiler/xla/client/lib/qr.h"
+#include "tensorflow/compiler/xla/client/lib/triangular_solve.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/client/xla_computation.h"
#include "tensorflow/compiler/xla/executable_run_options.h"
@@ -644,6 +647,15 @@
return xla::ConstantLiteral(&builder_, literal);
}
+LocalOp LocalComputationBuilder::Iota(PrimitiveType element_type, int64 size) {
+ return xla::Iota(&builder_, element_type, size);
+}
+
+LocalOp LocalComputationBuilder::BroadcastedIota(const Shape& shape,
+ int64 dimension) {
+ return xla::Iota(&builder_, shape, dimension);
+}
+
LocalOp LocalComputationBuilder::Broadcast(
const LocalOp& operand, absl::Span<const int64> broadcast_sizes) {
return xla::Broadcast(operand.op(), broadcast_sizes);
@@ -780,6 +792,21 @@
return xla::Call(&builder_, local_computation.computation(), xla_ops);
}
+LocalOp LocalComputationBuilder::CustomCall(
+ const string& call_target_name, absl::Span<const LocalOp> operands,
+ const Shape& shape_with_layout,
+ const std::vector<Shape>& operand_shapes_with_layout,
+ const string& opaque) {
+ std::vector<XlaOp> xla_ops;
+ xla_ops.reserve(operands.size());
+ for (const auto& op : operands) {
+ xla_ops.push_back(op.op());
+ }
+ return xla::CustomCallWithLayout(&builder_, call_target_name, xla_ops,
+ shape_with_layout,
+ operand_shapes_with_layout, opaque);
+}
+
LocalOp LocalComputationBuilder::Transpose(
const LocalOp& operand, absl::Span<const int64> permutation) {
return xla::Transpose(operand.op(), permutation);
@@ -865,6 +892,27 @@
return xla::Sort(keys.op(), {values.op()}, dimension);
}
+LocalOp LocalComputationBuilder::Cholesky(const LocalOp& a) {
+ return xla::Cholesky(a.op());
+}
+
+LocalOp LocalComputationBuilder::QR(const LocalOp& a, bool full_matrices) {
+ XlaBuilder* builder = a.op().builder();
+ return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> {
+ TF_ASSIGN_OR_RETURN(auto qr, xla::QRDecomposition(a.op(), full_matrices));
+ return xla::Tuple(builder, {qr.q, qr.r});
+ });
+}
+
+LocalOp LocalComputationBuilder::TriangularSolve(const LocalOp& a,
+ const LocalOp& b,
+ bool left_side, bool lower,
+ bool transpose_a,
+ bool conjugate_a) {
+ return xla::TriangularSolve(a.op(), b.op(), left_side, lower, transpose_a,
+ conjugate_a);
+}
+
StatusOr<LocalComputation*> LocalComputationBuilder::BuildConstantSubGraph(
const LocalOp& operand) {
TF_ASSIGN_OR_RETURN(XlaComputation computation,
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h
index 149e445..c6e58ac 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.h
+++ b/tensorflow/compiler/xla/python/local_computation_builder.h
@@ -286,6 +286,10 @@
LocalOp ConstantLiteral(const Literal& literal);
+ LocalOp Iota(PrimitiveType element_type, int64 size);
+
+ LocalOp BroadcastedIota(const Shape& shape, int64 dimension);
+
LocalOp Broadcast(const LocalOp& operand,
absl::Span<const int64> broadcast_sizes);
@@ -352,6 +356,12 @@
LocalOp Call(const LocalComputation& local_computation,
absl::Span<const LocalOp> operands);
+ LocalOp CustomCall(const string& call_target_name,
+ absl::Span<const LocalOp> operands,
+ const Shape& shape_with_layout,
+ const std::vector<Shape>& operand_shapes_with_layout,
+ const string& opaque);
+
LocalOp Transpose(const LocalOp& operand,
absl::Span<const int64> permutation);
@@ -394,6 +404,13 @@
LocalOp SortKeyVal(const LocalOp& keys, const LocalOp& values,
int64 dimension);
+ LocalOp QR(const LocalOp& a, bool full_matrices);
+
+ LocalOp Cholesky(const LocalOp& a);
+
+ LocalOp TriangularSolve(const LocalOp& a, const LocalOp& b, bool left_side,
+ bool lower, bool transpose_a, bool conjugate_a);
+
StatusOr<LocalComputation*> BuildConstantSubGraph(const LocalOp& operand);
#define _FORWARD(method_name, return_sig, args_sig) \
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i
index d23d693..11fb00e 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.i
+++ b/tensorflow/compiler/xla/python/local_computation_builder.i
@@ -1051,6 +1051,8 @@
%unignore xla::swig::LocalComputationBuilder::Outfeed;
%unignore xla::swig::LocalComputationBuilder::ConstantLiteral;
%unignore xla::swig::LocalComputationBuilder::ConstantR0;
+%unignore xla::swig::LocalComputationBuilder::Iota;
+%unignore xla::swig::LocalComputationBuilder::BroadcastedIota;
%unignore xla::swig::LocalComputationBuilder::Broadcast;
%unignore xla::swig::LocalComputationBuilder::BroadcastInDim;
%unignore xla::swig::LocalComputationBuilder::Pad;
@@ -1144,6 +1146,10 @@
%unignore xla::swig::LocalComputationBuilder::Imag;
%unignore xla::swig::LocalComputationBuilder::Conj;
%unignore xla::swig::LocalComputationBuilder::Complex;
+%unignore xla::swig::LocalComputationBuilder::Cholesky;
+%unignore xla::swig::LocalComputationBuilder::QR;
+%unignore xla::swig::LocalComputationBuilder::TriangularSolve;
+%unignore xla::swig::LocalComputationBuilder::CustomCall;
%unignore xla::swig::DeleteLocalComputation;
%unignore xla::swig::DestructureLocalShapedBufferTuple;
%unignore xla::swig::DestructureXrtAllocationTuple;
diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py
index c91a2aa..4166fa0 100644
--- a/tensorflow/compiler/xla/python/xla_client.py
+++ b/tensorflow/compiler/xla/python/xla_client.py
@@ -831,6 +831,33 @@
return self.ParameterWithShape(
Shape.from_pyval(value), name=name, parameter_num=parameter_num)
+ def Iota(self, dtype, size):
+ """Enqueues an iota constant onto the computation.
+
+ Args:
+ dtype: expected numpy dtype of the output.
+ size: integer, the number of elements in the array.
+
+ Returns:
+ A LocalOp representing the added iota constant.
+ """
+ element_type = DTYPE_TO_XLA_ELEMENT_TYPE[str(np.dtype(dtype))]
+ return self._client.Iota(element_type, size)
+
+ def BroadcastedIota(self, dtype, shape, dimension):
+ """Enqueues a broadcasted iota constant onto the computation.
+
+ Args:
+ dtype: expected numpy dtype of the output.
+ shape: tuple of integers, the expected output shape (dimensions).
+ dimension: positive integer, dimension along which to increment values.
+
+ Returns:
+ A LocalOp representing the added broadcasted iota constant.
+ """
+ xla_shape = Shape.array_shape(dtype, shape)
+ return self._client.BroadcastedIota(xla_shape, dimension)
+
def Broadcast(self, operand, sizes):
"""Enqueues a broadcast operation onto the computation.
@@ -1102,6 +1129,31 @@
"""
return self._client.Call(computation_to_apply.computation, operands)
+ def CustomCall(self,
+ call_target_name,
+ operands,
+ shape_with_layout,
+ operand_shapes_with_layout,
+ opaque=None):
+ """Enqueues a custom call operation onto the computation.
+
+ Args:
+ call_target_name: the name of the function to call.
+ operands: an iterable of LocalOp. The number and types of operands must
+ match the arity of `operand_shapes_with_layout`.
+ shape_with_layout: the shape of the operator's output, with layout.
+ operand_shapes_with_layout: the shapes of `operands`, including the
+ expected layouts.
+ opaque: an opaque string passed to the backend.
+
+ Returns:
+ A LocalOp representing the added custom call op.
+ """
+ opaque = opaque or ''
+ return self._client.CustomCall(call_target_name, operands,
+ shape_with_layout,
+ operand_shapes_with_layout, opaque)
+
def Map(self, operands, computation_to_apply, dimensions):
"""Enqueues a map operation onto the computation.
@@ -1411,6 +1463,20 @@
"""Enqueues a key-value sort operation onto the computation."""
return self._client.SortKeyVal(keys, values, dimension)
+ def Cholesky(self, a):
+ """Enqueues a Cholesky decomposition onto the computation."""
+ return self._client.Cholesky(a)
+
+ def QR(self, a, full_matrices=True):
+ """Enqueues a QR decomposition onto the computation."""
+ return self._client.QR(a, full_matrices)
+
+ def TriangularSolve(self, a, b, left_side=False, lower=False,
+ transpose_a=False, conjugate_a=False):
+ """Enqueues a triangular-solve operation onto the computation."""
+ return self._client.TriangularSolve(
+ a, b, left_side, lower, transpose_a, conjugate_a)
+
def _forward_methods_to_local_builder():
"""Forward remaining ComputationBuilder methods to the C API.
diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py
index 21b5c93..95c6dc8 100644
--- a/tensorflow/compiler/xla/python/xla_client_test.py
+++ b/tensorflow/compiler/xla/python/xla_client_test.py
@@ -18,6 +18,7 @@
from __future__ import division
from __future__ import print_function
+import functools
import itertools
import threading
@@ -51,9 +52,11 @@
def _ExecuteAndCompareExact(self, c, arguments=(), expected=None):
self._ExecuteAndAssertWith(np.testing.assert_equal, c, arguments, expected)
- def _ExecuteAndCompareClose(self, c, arguments=(), expected=None):
- self._ExecuteAndAssertWith(np.testing.assert_allclose, c, arguments,
- expected)
+ def _ExecuteAndCompareClose(self, c, arguments=(), expected=None, rtol=1e-7,
+ atol=0):
+ self._ExecuteAndAssertWith(
+ functools.partial(np.testing.assert_allclose, rtol=rtol, atol=atol),
+ c, arguments, expected)
def NumpyArrayF32(*args, **kwargs):
@@ -143,6 +146,17 @@
c.Pow(c.Constant(NumpyArrayF64([1.5, 2.5, 3.0])), c.ConstantF64Scalar(2.))
self._ExecuteAndCompareClose(c, expected=[2.25, 6.25, 9.])
+ def testIota(self):
+ c = self._NewComputation()
+ c.Iota(np.float32, 10)
+ self._ExecuteAndCompareExact(c, expected=np.arange(10, dtype=np.float32))
+
+ def testBroadcastedIota(self):
+ c = self._NewComputation()
+ c.BroadcastedIota(np.int64, (2, 3), 1)
+ expected = np.array([[0, 1, 2], [0, 1, 2]], dtype=np.int64)
+ self._ExecuteAndCompareExact(c, expected=expected)
+
def testBooleanAnd(self):
c = self._NewComputation()
c.And(
@@ -1057,6 +1071,38 @@
self.assertTrue(np.all(lo <= result))
self.assertTrue(np.all(result < hi))
+ def testCholesky(self):
+ l = np.array([[4, 0, 0, 0], [6, 5, 0, 0], [2, 14, 16, 0], [3, 6, 1, 4]],
+ dtype=np.float32)
+ c = self._NewComputation()
+ c.Cholesky(c.Constant(np.dot(l, l.T)))
+ self._ExecuteAndCompareClose(c, expected=l, rtol=1e-4)
+
+ def testQR(self):
+ a = np.array(
+ [[4, 6, 8, 10], [6, 45, 54, 63], [8, 54, 146, 166], [10, 63, 166, 310]],
+ dtype=np.float32)
+ c = self._NewComputation()
+ c.QR(c.Constant(a), full_matrices=True)
+ q, r = self._Execute(c, ())
+ np.testing.assert_allclose(np.dot(q, r), a, rtol=1e-4)
+
+ def testTriangularSolve(self):
+ a_vals = np.array(
+ [[2, 0, 0, 0], [3, 6, 0, 0], [4, 7, 9, 0], [5, 8, 10, 11]],
+ dtype=np.float32)
+ b_vals = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
+ dtype=np.float32)
+
+ c = self._NewComputation()
+ c.TriangularSolve(c.Constant(a_vals), c.Constant(b_vals), left_side=False,
+ lower=True, transpose_a=True)
+ self._ExecuteAndCompareClose(c, expected=np.array([
+ [0.5, 0.08333334, 0.04629629, 0.03367003],
+ [2.5, -0.25, -0.1388889, -0.1010101],
+ [4.5, -0.58333331, -0.32407406, -0.23569024],
+ ], dtype=np.float32), rtol=1e-4)
+
def testIsConstant(self):
c = self._NewComputation()
a = c.ConstantS32Scalar(3)
diff --git a/tensorflow/compiler/xla/reference_util.cc b/tensorflow/compiler/xla/reference_util.cc
index ceb5e74..a27e200 100644
--- a/tensorflow/compiler/xla/reference_util.cc
+++ b/tensorflow/compiler/xla/reference_util.cc
@@ -21,7 +21,6 @@
#include "absl/memory/memory.h"
#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/literal_util.h"
-#include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h"
#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/shape_inference.h"
@@ -32,48 +31,19 @@
namespace xla {
-namespace {
-
-template <typename T>
-std::unique_ptr<Array2D<T>> MatmulArray2DImpl(
- const Array2D<T>& lhs, const Array2D<T>& rhs,
- const std::function<void(
- const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, int64 n,
- int64 k, int32 transpose_lhs, int32 transpose_rhs)>& impl_fn) {
- CHECK_EQ(lhs.width(), rhs.height());
- int m = lhs.height();
- int n = rhs.width();
- int k = lhs.width();
- auto result = absl::make_unique<Array2D<T>>(m, n);
- // Because Eigen is a header-oriented library, make sure that the Eigen code
- // is the same as the code used by the CPU backend (otherwise the linker will
- // randomly pick *some* definition).
- impl_fn(
- /*run_options_ptr=*/nullptr, result->data(), rhs.data(), lhs.data(), n, m,
- k,
- /*transpose_lhs=*/0,
- /*transpose_rhs=*/0);
- return result;
-}
-
-} // namespace
-
/* static */ std::unique_ptr<Array2D<Eigen::half>> ReferenceUtil::MatmulArray2D(
const Array2D<Eigen::half>& lhs, const Array2D<Eigen::half>& rhs) {
- return MatmulArray2DImpl<Eigen::half>(
- lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF16);
+ return HloEvaluator::MatmulArray2D(lhs, rhs);
}
/* static */ std::unique_ptr<Array2D<float>> ReferenceUtil::MatmulArray2D(
const Array2D<float>& lhs, const Array2D<float>& rhs) {
- return MatmulArray2DImpl<float>(
- lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF32);
+ return HloEvaluator::MatmulArray2D(lhs, rhs);
}
/* static */ std::unique_ptr<Array2D<double>> ReferenceUtil::MatmulArray2D(
const Array2D<double>& lhs, const Array2D<double>& rhs) {
- return MatmulArray2DImpl<double>(
- lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF64);
+ return HloEvaluator::MatmulArray2D(lhs, rhs);
}
/* static */ std::unique_ptr<Array2D<double>> ReferenceUtil::Array2DF32ToF64(
diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 81e71ee..4aaa8a5 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -241,6 +241,7 @@
":hlo_casting_utils",
":hlo_query",
":shape_inference",
+ "//tensorflow/compiler/xla:array2d",
"//tensorflow/compiler/xla:literal",
"//tensorflow/compiler/xla:literal_util",
"//tensorflow/compiler/xla:shape_util",
@@ -249,6 +250,7 @@
"//tensorflow/compiler/xla:util",
"//tensorflow/compiler/xla:window_util",
"//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_matmul",
"//tensorflow/core:lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base",
@@ -1910,6 +1912,41 @@
],
)
+cc_library(
+ name = "dynamic_dimension_inference",
+ srcs = ["dynamic_dimension_inference.cc"],
+ hdrs = ["dynamic_dimension_inference.h"],
+ deps = [
+ ":hlo",
+ "//tensorflow/compiler/xla:status",
+ "//tensorflow/compiler/xla:statusor",
+ "//tensorflow/compiler/xla:types",
+ "//tensorflow/core:lib",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/types:span",
+ ],
+)
+
+tf_cc_test(
+ name = "dynamic_dimension_inference_test",
+ srcs = ["dynamic_dimension_inference_test.cc"],
+ deps = [
+ ":dynamic_dimension_inference",
+ "//tensorflow/compiler/xla:debug_options_flags",
+ "//tensorflow/compiler/xla:literal",
+ "//tensorflow/compiler/xla:shape_util",
+ "//tensorflow/compiler/xla:test",
+ "//tensorflow/compiler/xla:test_helpers",
+ "//tensorflow/compiler/xla:xla_data_proto",
+ "//tensorflow/compiler/xla/client:xla_builder",
+ "//tensorflow/compiler/xla/service:hlo",
+ "//tensorflow/compiler/xla/service:hlo_matchers",
+ "//tensorflow/compiler/xla/service:hlo_runner",
+ "//tensorflow/compiler/xla/tests:hlo_test_base",
+ "//tensorflow/core:test",
+ ],
+)
+
tf_cc_test(
name = "reshape_mover_test",
srcs = ["reshape_mover_test.cc"],
@@ -2067,7 +2104,8 @@
srcs = ["hlo_computation_test.cc"],
deps = [
":hlo",
- ":hlo_matchers",
+ ":pattern_matcher",
+ ":pattern_matcher_gmock",
"//tensorflow/compiler/xla:literal",
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/compiler/xla:test",
@@ -2661,7 +2699,6 @@
":algebraic_simplifier",
":computation_layout",
":hlo",
- ":hlo_matchers",
":layout_assignment",
":pattern_matcher",
":pattern_matcher_gmock",
@@ -2675,6 +2712,7 @@
"//tensorflow/compiler/xla/service:hlo_parser",
"//tensorflow/compiler/xla/tests:hlo_test_base",
"//tensorflow/compiler/xla/tests:test_utils",
+ "//tensorflow/compiler/xla/tests:xla_internal_test_main",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"@com_google_absl//absl/types:span",
@@ -3127,6 +3165,7 @@
name = "hlo_graph_dumper",
srcs = [
"hlo_graph_dumper.cc",
+ "hlo_graph_html_renderer.cc",
],
hdrs = ["hlo_graph_dumper.h"],
deps = [
diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner.cc b/tensorflow/compiler/xla/service/ar_crs_combiner.cc
index c11452a..47d2c7e 100644
--- a/tensorflow/compiler/xla/service/ar_crs_combiner.cc
+++ b/tensorflow/compiler/xla/service/ar_crs_combiner.cc
@@ -36,31 +36,47 @@
namespace m = match;
-// If the argument instruction is a CRS in the sequence
-// AR -> Convert -> Add -> CRS
-// then return the AR in the sequence.
-// TODO(b/117554291): Rewrite this to recognize more general patterns,
-// not just the specific one of AR -> Add -> Convert -> CRS.
-absl::optional<HloInstruction*> MatchesArCrsPattern(
- HloInstruction* instruction) {
- HloInstruction *ar, *convert, *add, *crs;
- if (Match(instruction,
- m::CrossReplicaSum(
- &crs, m::Add(&add, m::Op(),
- m::Convert(&convert,
- m::CrossReplicaSum(&ar, m::Op()))))) &&
- ar->users().size() == 1 && ar->shape().element_type() == BF16 &&
- convert->shape().element_type() == F32 && !crs->all_reduce_id()) {
- return ar;
+// Returns true iff the argument instruction is an AllReduce, followed by a
+// certain sequence of instructions and then a CRS. It must be possible to move
+// the AR past each instruction in the sequence.
+bool MatchesArCrsPattern(HloInstruction* instruction) {
+ auto can_ar_move_past_instruction = [](HloInstruction* instruction) -> bool {
+ if (instruction->user_count() != 1) {
+ return false;
+ }
+ auto opcode = instruction->opcode();
+ return opcode == HloOpcode::kBitcast || opcode == HloOpcode::kTranspose ||
+ opcode == HloOpcode::kReshape || opcode == HloOpcode::kConvert ||
+ opcode == HloOpcode::kAdd || opcode == HloOpcode::kSubtract ||
+ opcode == HloOpcode::kMultiply;
+ };
+
+ auto computation_is_addition = [](HloComputation* c) {
+ return c->instruction_count() == 3 &&
+ Match(c->root_instruction(), m::Add(m::Parameter(), m::Parameter()));
+ };
+
+ if (!instruction->IsCrossModuleAllReduce() ||
+ !computation_is_addition(instruction->called_computations()[0]) ||
+ instruction->user_count() != 1) {
+ return false;
}
- return absl::optional<HloInstruction*>();
+ auto next = instruction->users()[0];
+ while (!next->IsCrossReplicaAllReduce()) {
+ if (can_ar_move_past_instruction(next)) {
+ next = next->users()[0];
+ } else {
+ return false;
+ }
+ }
+ return computation_is_addition(next->called_computations()[0]);
}
} // namespace
absl::optional<HloInstruction*> ArCrsCombiner::WhileFromBodyParameter(
HloInstruction* instruction) {
- CHECK(HloOpcode::kParameter == instruction->opcode());
+ CHECK_EQ(HloOpcode::kParameter, instruction->opcode());
HloComputation* computation = instruction->parent();
auto caller_instructions = call_graph_->GetComputationCallers(computation);
if (caller_instructions.size() == 1) {
@@ -120,7 +136,7 @@
return false;
}
for (auto tuple : tuples) {
- CHECK(tuple->opcode() == HloOpcode::kTuple);
+ CHECK_EQ(tuple->opcode(), HloOpcode::kTuple);
if (!InstructionsComputeSameValue(tuple->mutable_operand(i1),
tuple->mutable_operand(i2),
visited_pairs)) {
@@ -160,13 +176,6 @@
if (opcode1 != i2->opcode() || operands1.size() != i2->operands().size()) {
return false;
}
- if (opcode1 == HloOpcode::kConstant || i1->IsCrossModuleAllReduce()) {
- return i1->Identical(
- *i2,
- /*eq_operands=*/std::equal_to<const HloInstruction*>(),
- /*eq_computations=*/std::equal_to<const HloComputation*>(),
- /*layout_sensitive=*/false);
- }
visited_pairs->emplace(min_uid, max_uid);
for (int i = 0; i < operands1.size(); ++i) {
auto operand1 = operands1[i];
@@ -175,22 +184,35 @@
return false;
}
}
+ if (opcode1 == HloOpcode::kParameter) {
+ // In the general case, we don't try to prove equality of parameters.
+ // We only try in the context of get-tuple-element
+ // (see TupleElementsComputeSameValue).
+ return false;
+ }
if (opcode1 == HloOpcode::kGetTupleElement) {
- if (i1->tuple_index() == i2->tuple_index()) {
- return true;
- }
- return TupleElementsComputeSameValue(operands1[0], i1->tuple_index(),
+ return i1->tuple_index() == i2->tuple_index() ||
+ TupleElementsComputeSameValue(operands1[0], i1->tuple_index(),
i2->tuple_index(), visited_pairs);
}
- return true;
+ // Don't check that the operands are identical, because Identical can
+ // return false for instructions that compute the same value but are not
+ // identical, which we don't want. We have checked the arguments with
+ // InstructionsComputeSameValue earlier.
+ auto eq_instructions = [](const HloInstruction* i1,
+ const HloInstruction* i2) -> bool { return true; };
+ auto eq_computations = [](const HloComputation* a, const HloComputation* b) {
+ return *a == *b;
+ };
+ return i1->Identical(*i2, eq_instructions, eq_computations,
+ /*layout_sensitive=*/false);
}
void ArCrsCombiner::GroupAllReducesById(HloModule* module) {
for (HloComputation* computation : module->MakeNonfusionComputations()) {
for (HloInstruction* instruction : computation->instructions()) {
- auto ar = MatchesArCrsPattern(instruction);
- if (ar) {
- all_reduce_map_[*((*ar)->all_reduce_id())].push_back(*ar);
+ if (MatchesArCrsPattern(instruction)) {
+ all_reduce_map_[*(instruction->all_reduce_id())].push_back(instruction);
}
}
}
@@ -198,21 +220,23 @@
void ArCrsCombiner::KeepProvablyEqualInstructionGroups() {
for (auto it : all_reduce_map_) {
+ auto all_reduce_id = it.first;
auto instruction_vec = it.second;
CHECK_EQ(instruction_vec.size(), num_spatial_partitions_);
-
auto instr_0 = instruction_vec[0];
- auto add_0 = instr_0->users()[0]->users()[0];
- CHECK(HloOpcode::kAdd == add_0->opcode());
-
for (int i = 1; i < instruction_vec.size(); ++i) {
auto instr_i = instruction_vec[i];
- auto add_i = instr_i->users()[0]->users()[0];
- CHECK(HloOpcode::kAdd == add_i->opcode());
+ auto next_0 = instr_0->users()[0];
+ auto next_i = instr_i->users()[0];
absl::flat_hash_map<int64, int64> visited_pairs;
- if (!InstructionsComputeSameValue(add_0, add_i, &visited_pairs)) {
- all_reduce_map_.erase(it.first);
- }
+ do {
+ if (!InstructionsComputeSameValue(next_0, next_i, &visited_pairs)) {
+ all_reduce_map_.erase(all_reduce_id);
+ break;
+ }
+ next_0 = next_0->users()[0];
+ next_i = next_i->users()[0];
+ } while (!next_0->IsCrossReplicaAllReduce());
}
}
}
@@ -221,55 +245,51 @@
if (all_reduce_map_.empty()) {
return false;
}
-
- auto computation_is_addition = [](HloComputation* c) {
- return c->instruction_count() == 3 &&
- Match(c->root_instruction(), m::Add(m::Parameter(), m::Parameter()));
- };
-
for (auto it : all_reduce_map_) {
auto instruction_vec = it.second;
for (auto all_reduce : instruction_vec) {
auto parent_computation = all_reduce->parent();
- auto convert = all_reduce->users()[0];
- auto add = convert->users()[0];
- auto crs = add->users()[0];
-
- if (!computation_is_addition(all_reduce->called_computations()[0]) ||
- !computation_is_addition(crs->called_computations()[0])) {
- continue;
- }
- HloInstruction* other_summand = (add->operands()[0] == convert)
- ? add->operands()[1]
- : add->operands()[0];
- // Remove the AllReduce and replace the CRS with:
- // AllReduce - (other_summand * (num_spatial_partitions_ - 1))
- TF_CHECK_OK(
- all_reduce->ReplaceAllUsesWith(all_reduce->mutable_operand(0)));
- crs->set_all_reduce_id(all_reduce->all_reduce_id());
- auto new_shape = crs->shape();
- HloInstruction* to_subtract;
- if (num_spatial_partitions_ == 2) {
- to_subtract = other_summand;
- } else {
- Literal partitions_minus_1_lit = Literal(new_shape);
- partitions_minus_1_lit.PopulateWithValue<float>(
- num_spatial_partitions_ - 1);
- auto partitions_minus_1_const = parent_computation->AddInstruction(
- HloInstruction::CreateConstant(partitions_minus_1_lit.Clone()));
- to_subtract =
- parent_computation->AddInstruction(HloInstruction::CreateBinary(
- new_shape, HloOpcode::kMultiply, other_summand,
- partitions_minus_1_const));
- }
- auto sub =
- parent_computation->AddInstruction(HloInstruction::CreateBinary(
- new_shape, HloOpcode::kSubtract, crs, to_subtract));
- TF_CHECK_OK(crs->ReplaceAllUsesWith(sub));
+ auto all_reduce_id = all_reduce->all_reduce_id();
+ auto prev = all_reduce->mutable_operand(0);
+ auto next = all_reduce->users()[0];
+ TF_CHECK_OK(all_reduce->ReplaceUseWith(next, prev));
TF_CHECK_OK(parent_computation->RemoveInstruction(all_reduce));
+ while (!next->IsCrossReplicaAllReduce()) {
+ switch (next->opcode()) {
+ case HloOpcode::kBitcast:
+ case HloOpcode::kTranspose:
+ case HloOpcode::kReshape:
+ case HloOpcode::kConvert:
+ case HloOpcode::kMultiply:
+ break;
+ case HloOpcode::kAdd:
+ case HloOpcode::kSubtract: {
+ auto other_operand = (next->operands()[0] == prev)
+ ? next->operands()[1]
+ : next->operands()[0];
+ // To move the AR past the addition/subtraction, we need to divide
+ // other_operand by the number of spatial partitions.
+ auto shape = other_operand->shape();
+ Literal lit(shape);
+ lit.PopulateWithValue<float>(num_spatial_partitions_);
+ auto divisor = parent_computation->AddInstruction(
+ HloInstruction::CreateConstant(lit.Clone()));
+ auto division =
+ parent_computation->AddInstruction(HloInstruction::CreateBinary(
+ shape, HloOpcode::kDivide, other_operand, divisor));
+ TF_CHECK_OK(other_operand->ReplaceUseWith(next, division));
+ break;
+ }
+ default:
+ LOG(FATAL) << "Unexpected instruction: " << next->ToShortString();
+ }
+ prev = next;
+ next = next->users()[0];
+ }
+ // The AllReduce and the CRS are combined to an all-core AllReduce.
+ next->set_all_reduce_id(all_reduce_id);
}
}
-
return true;
}
diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner.h b/tensorflow/compiler/xla/service/ar_crs_combiner.h
index f6a7ef7..6be7e10 100644
--- a/tensorflow/compiler/xla/service/ar_crs_combiner.h
+++ b/tensorflow/compiler/xla/service/ar_crs_combiner.h
@@ -25,9 +25,12 @@
namespace xla {
-// Combine an AllReduce and a CrossReplicaSum when they are close to each other
-// in the graph, to use an efficient CrossReplicaSum implementation that
-// fully utilizes the interconnect bandwidth.
+// When the HLO graph contains an AllReduce, followed by some simple linear
+// operations, followed by a CrossReplicaSum, we can combine the AR and the CRS,
+// to use an efficient CrossReplicaSum implementation that fully utilizes the
+// interconnect bandwidth.
+// Such sequences appear in spatially partitioned models.
+// This pass must run right after spatial partitioning.
class ArCrsCombiner : public HloModulePass {
public:
ArCrsCombiner(int num_spatial_partitions)
diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc b/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc
index 9d5eaf6..2f7a53b 100644
--- a/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc
+++ b/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc
@@ -48,6 +48,43 @@
EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
}
+TEST_F(ArCrsCombinerTest, SameValueTestBasecase2) {
+ const char* module_str = R"(
+HloModule foobar
+
+ENTRY %entrycomp (x: f32[]) -> (f32[], f32[]) {
+ %x = f32[] parameter(0)
+ ROOT %tuple = (f32[], f32[]) tuple(%x, %x)
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto root_tuple = module->entry_computation()->root_instruction();
+ auto i1 = root_tuple->operands()[0];
+ auto i2 = root_tuple->operands()[1];
+ EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
+}
+
+TEST_F(ArCrsCombinerTest, SameValueTestBasecase3) {
+ const char* module_str = R"(
+HloModule foobar
+
+ENTRY %entrycomp (x: f32[], y: f32[]) -> (f32[], f32[]) {
+ %x = f32[] parameter(0)
+ %y = f32[] parameter(1)
+ ROOT %tuple = (f32[], f32[]) tuple(%x, %y)
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto root_tuple = module->entry_computation()->root_instruction();
+ auto i1 = root_tuple->operands()[0];
+ auto i2 = root_tuple->operands()[1];
+ EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
+}
+
TEST_F(ArCrsCombinerTest, SameValueTestNumOperands) {
const char* module_str = R"(
HloModule foobar
@@ -69,6 +106,46 @@
EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
}
+TEST_F(ArCrsCombinerTest, SameValueTestSliceIndicesMatch) {
+ const char* module_str = R"(
+HloModule foobar
+
+ENTRY %entrycomp (p: f32[2]) -> (f32[1], f32[1]) {
+ %p = f32[2] parameter(0)
+ %slice.1 = f32[1] slice(f32[2] %p), slice={[0:1]}
+ %slice.2 = f32[1] slice(f32[2] %p), slice={[0:1]}
+ ROOT %tuple = (f32[1], f32[1]) tuple(%slice.1, %slice.2)
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto root_tuple = module->entry_computation()->root_instruction();
+ auto i1 = root_tuple->operands()[0];
+ auto i2 = root_tuple->operands()[1];
+ EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
+}
+
+TEST_F(ArCrsCombinerTest, SameValueTestSliceIndicesDontMatch) {
+ const char* module_str = R"(
+HloModule foobar
+
+ENTRY %entrycomp (p: f32[2]) -> (f32[1], f32[1]) {
+ %p = f32[2] parameter(0)
+ %slice.1 = f32[1] slice(f32[2] %p), slice={[0:1]}
+ %slice.2 = f32[1] slice(f32[2] %p), slice={[1:2]}
+ ROOT %tuple = (f32[1], f32[1]) tuple(%slice.1, %slice.2)
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto root_tuple = module->entry_computation()->root_instruction();
+ auto i1 = root_tuple->operands()[0];
+ auto i2 = root_tuple->operands()[1];
+ EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
+}
+
TEST_F(ArCrsCombinerTest, SameValueTestTupleElementSameIndex) {
const char* module_str = R"(
HloModule foobar
@@ -249,11 +326,27 @@
EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2));
}
-TEST_F(ArCrsCombinerTest, RewritePatternArConvertAddCrs) {
+void CompareReplicaGroups(const std::vector<ReplicaGroup>& groups_before,
+ const std::vector<ReplicaGroup>& groups_after) {
+ ASSERT_EQ(groups_before.size(), groups_after.size());
+ for (int i = 0; i < groups_before.size(); ++i) {
+ // Somewhat verbose way to compare the replica_ids, because EqualsProto
+ // is not available in the open-source build.
+ auto group_before = groups_before[i];
+ std::vector<int64> ids_before(group_before.replica_ids().begin(),
+ group_before.replica_ids().end());
+ auto group_after = groups_after[i];
+ std::vector<int64> ids_after(group_after.replica_ids().begin(),
+ group_after.replica_ids().end());
+ EXPECT_EQ(ids_before, ids_after);
+ }
+}
+
+TEST_F(ArCrsCombinerTest, RewriteArConvertCrs) {
const char* module_str = R"(
HloModule foobar
-%binary_add (a: bf16[], b: bf16[]) -> bf16[] {
+%sum.bf16 (a: bf16[], b: bf16[]) -> bf16[] {
%a = bf16[] parameter(0)
%b = bf16[] parameter(1)
ROOT %add = bf16[] add(%a, %b)
@@ -265,48 +358,40 @@
ROOT %add = f32[] add(%x, %y)
}
-ENTRY %entrycomp (p: f32[2,2]) -> (f32[2,2], f32[2,2]) {
- %p = f32[2,2] parameter(0)
- %constant.bf16 = bf16[2,2] constant(bf16[2,2] {{1, 2}, {3, 4}})
- %constant.f32 = f32[2,2] constant(f32[2,2] {{1, 2}, {3, 4}})
+ENTRY %entrycomp (p: bf16[]) -> (f32[], f32[]) {
+ %p = bf16[] parameter(0)
- %cross-replica-sum.ar.1 = bf16[2,2]
- cross-replica-sum(%constant.bf16),
+ %cross-replica-sum.ar.1 = bf16[]
+ cross-replica-sum(%p),
replica_groups={{0},{1}},
all_reduce_id=1,
- to_apply=%binary_add,
+ to_apply=%sum.bf16,
sharding={maximal device=0}
- %convert.1 = f32[2,2]
+ %convert.1 = f32[]
convert(%cross-replica-sum.ar.1),
sharding={maximal device=0}
- %add.1 = f32[2,2]
- add(%constant.f32, %convert.1),
- sharding={maximal device=0}
- %cross-replica-sum.1 = f32[2,2]
- cross-replica-sum(%add.1),
+ %cross-replica-sum.1 = f32[]
+ cross-replica-sum(%convert.1),
replica_groups={{0,1}},
to_apply=%sum.f32,
sharding={maximal device=0}
- %cross-replica-sum.ar.2 = bf16[2,2]
- cross-replica-sum(%constant.bf16),
+ %cross-replica-sum.ar.2 = bf16[]
+ cross-replica-sum(%p),
replica_groups={{0},{1}},
all_reduce_id=1,
- to_apply=%binary_add,
+ to_apply=%sum.bf16,
sharding={maximal device=1}
- %convert.2 = f32[2,2]
+ %convert.2 = f32[]
convert(%cross-replica-sum.ar.2),
sharding={maximal device=1}
- %add.2 = f32[2,2]
- add(%constant.f32, %convert.2),
- sharding={maximal device=1}
- %cross-replica-sum.2 = f32[2,2]
- cross-replica-sum(%add.2),
+ %cross-replica-sum.2 = f32[]
+ cross-replica-sum(%convert.2),
replica_groups={{0,1}},
to_apply=%sum.f32,
sharding={maximal device=1}
- ROOT %tuple = (f32[2,2], f32[2,2])
+ ROOT %tuple = (f32[], f32[])
tuple(%cross-replica-sum.1, %cross-replica-sum.2),
sharding={{maximal device=0}, {maximal device=1}}
}
@@ -321,30 +406,156 @@
auto changed = combiner.Run(module.get()).ValueOrDie();
EXPECT_TRUE(changed);
EXPECT_THAT(module->entry_computation()->root_instruction(),
- op::Tuple(op::Subtract(op::CrossReplicaSum(), op::Constant()),
- op::Subtract(op::CrossReplicaSum(), op::Constant())));
- auto sub = module->entry_computation()->root_instruction()->operands()[0];
- auto crs_after = sub->operands()[0];
+ op::Tuple(op::CrossReplicaSum(op::Convert(op::Parameter())),
+ op::CrossReplicaSum(op::Convert(op::Parameter()))));
+ auto crs_after =
+ module->entry_computation()->root_instruction()->operands()[0];
auto replica_groups_after = crs_after->replica_groups();
- ASSERT_EQ(replica_groups_before.size(), replica_groups_after.size());
- for (int i = 0; i < replica_groups_before.size(); ++i) {
- // Somewhat verbose way to compare the replica_ids, because EqualsProto
- // is not available in the open-source build.
- auto group_before = replica_groups_before[i];
- std::vector<int64> ids_before(group_before.replica_ids().begin(),
- group_before.replica_ids().end());
- auto group_after = replica_groups_after[i];
- std::vector<int64> ids_after(group_after.replica_ids().begin(),
- group_after.replica_ids().end());
- EXPECT_EQ(ids_before, ids_after);
- }
+ CompareReplicaGroups(replica_groups_before, replica_groups_after);
}
-TEST_F(ArCrsCombinerTest, OtherSummandNotTheSameDontRewrite) {
+TEST_F(ArCrsCombinerTest, RewriteArBitcastCrs) {
const char* module_str = R"(
HloModule foobar
-%binary_add (a: bf16[], b: bf16[]) -> bf16[] {
+%sum.1 (a: f32[2,1], b: f32[2,1]) -> f32[2,1] {
+ %a = f32[2,1] parameter(0)
+ %b = f32[2,1] parameter(1)
+ ROOT %add = f32[2,1] add(%a, %b)
+}
+
+%sum.2 (x: f32[2], y: f32[2]) -> f32[2] {
+ %x = f32[2] parameter(0)
+ %y = f32[2] parameter(1)
+ ROOT %add = f32[2] add(%x, %y)
+}
+
+ENTRY %entrycomp (p: f32[2,1]) -> (f32[2], f32[2]) {
+ %p = f32[2,1] parameter(0)
+
+ %cross-replica-sum.ar.1 = f32[2,1]
+ cross-replica-sum(%p),
+ replica_groups={{0},{1}},
+ all_reduce_id=1,
+ to_apply=%sum.1,
+ sharding={maximal device=0}
+ %bitcast.1 = f32[2]{0} bitcast(f32[2,1]{1,0} %cross-replica-sum.ar.1)
+ %cross-replica-sum.1 = f32[2]
+ cross-replica-sum(%bitcast.1),
+ replica_groups={{0,1}},
+ to_apply=%sum.2,
+ sharding={maximal device=0}
+
+ %cross-replica-sum.ar.2 = f32[2,1]
+ cross-replica-sum(%p),
+ replica_groups={{0},{1}},
+ all_reduce_id=1,
+ to_apply=%sum.1,
+ sharding={maximal device=1}
+ %bitcast.2 = f32[2]{0} bitcast(f32[2,1]{1,0} %cross-replica-sum.ar.2)
+ %cross-replica-sum.2 = f32[2]
+ cross-replica-sum(%bitcast.2),
+ replica_groups={{0,1}},
+ to_apply=%sum.2,
+ sharding={maximal device=1}
+
+ ROOT %tuple = (f32[], f32[])
+ tuple(%cross-replica-sum.1, %cross-replica-sum.2),
+ sharding={{maximal device=0}, {maximal device=1}}
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto crs_before =
+ module->entry_computation()->root_instruction()->operands()[0];
+ auto replica_groups_before = crs_before->replica_groups();
+ ArCrsCombiner combiner(2);
+ auto changed = combiner.Run(module.get()).ValueOrDie();
+ EXPECT_TRUE(changed);
+ EXPECT_THAT(module->entry_computation()->root_instruction(),
+ op::Tuple(op::CrossReplicaSum(op::Bitcast(op::Parameter())),
+ op::CrossReplicaSum(op::Bitcast(op::Parameter()))));
+ auto crs_after =
+ module->entry_computation()->root_instruction()->operands()[0];
+ auto replica_groups_after = crs_after->replica_groups();
+ CompareReplicaGroups(replica_groups_before, replica_groups_after);
+}
+
+TEST_F(ArCrsCombinerTest, RewriteArMultiplyCrs) {
+ const char* module_str = R"(
+HloModule foobar
+
+%sum.f32 (x: f32[], y: f32[]) -> f32[] {
+ %x = f32[] parameter(0)
+ %y = f32[] parameter(1)
+ ROOT %add = f32[] add(%x, %y)
+}
+
+ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) {
+ %p = f32[] parameter(0)
+ %constant.f32 = f32[] constant(123)
+
+ %cross-replica-sum.ar.1 = f32[]
+ cross-replica-sum(%p),
+ replica_groups={{0},{1}},
+ all_reduce_id=1,
+ to_apply=%sum.f32,
+ sharding={maximal device=0}
+ %multiply.1 = f32[]
+ multiply(%cross-replica-sum.ar.1, %constant.f32),
+ sharding={maximal device=0}
+ %cross-replica-sum.1 = f32[]
+ cross-replica-sum(%multiply.1),
+ replica_groups={{0,1}},
+ to_apply=%sum.f32,
+ sharding={maximal device=0}
+
+ %cross-replica-sum.ar.2 = f32[]
+ cross-replica-sum(%p),
+ replica_groups={{0},{1}},
+ all_reduce_id=1,
+ to_apply=%sum.f32,
+ sharding={maximal device=1}
+ %multiply.2 = f32[]
+ multiply(%cross-replica-sum.ar.2, %constant.f32),
+ sharding={maximal device=1}
+ %cross-replica-sum.2 = f32[]
+ cross-replica-sum(%multiply.2),
+ replica_groups={{0,1}},
+ to_apply=%sum.f32,
+ sharding={maximal device=1}
+
+ ROOT %tuple = (f32[], f32[])
+ tuple(%cross-replica-sum.1, %cross-replica-sum.2),
+ sharding={{maximal device=0}, {maximal device=1}}
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto crs_before =
+ module->entry_computation()->root_instruction()->operands()[0];
+ auto replica_groups_before = crs_before->replica_groups();
+ ArCrsCombiner combiner(2);
+ auto changed = combiner.Run(module.get()).ValueOrDie();
+ EXPECT_TRUE(changed);
+ EXPECT_THAT(
+ module->entry_computation()->root_instruction(),
+ op::Tuple(
+ op::CrossReplicaSum(op::Multiply(op::Parameter(), op::Constant())),
+ op::CrossReplicaSum(op::Multiply(op::Parameter(), op::Constant()))));
+ auto crs_after =
+ module->entry_computation()->root_instruction()->operands()[0];
+ auto replica_groups_after = crs_after->replica_groups();
+ CompareReplicaGroups(replica_groups_before, replica_groups_after);
+}
+
+TEST_F(ArCrsCombinerTest, RewriteArConvertAddCrs) {
+ const char* module_str = R"(
+HloModule foobar
+
+%sum.bf16 (a: bf16[], b: bf16[]) -> bf16[] {
%a = bf16[] parameter(0)
%b = bf16[] parameter(1)
ROOT %add = bf16[] add(%a, %b)
@@ -356,49 +567,133 @@
ROOT %add = f32[] add(%x, %y)
}
-ENTRY %entrycomp (p: f32[2,2]) -> (f32[2,2], f32[2,2]) {
- %p = f32[2,2] parameter(0)
- %constant.bf16 = bf16[2,2] constant(bf16[2,2] {{1, 2}, {3, 4}})
- %constant.f32.1 = f32[2,2] constant(f32[2,2] {{1, 2}, {3, 4}})
- %constant.f32.2 = f32[2,2] constant(f32[2,2] {{3, 4}, {5, 6}})
+ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) {
+ %p = f32[] parameter(0)
+ %constant.bf16 = bf16[] constant(1)
+ %constant.f32 = f32[] constant(2)
- %cross-replica-sum.ar.1 = bf16[2,2]
+ %cross-replica-sum.ar.1 = bf16[]
cross-replica-sum(%constant.bf16),
replica_groups={{0},{1}},
all_reduce_id=1,
- to_apply=%binary_add,
+ to_apply=%sum.bf16,
sharding={maximal device=0}
- %convert.1 = f32[2,2]
+ %convert.1 = f32[]
convert(%cross-replica-sum.ar.1),
sharding={maximal device=0}
- %add.1 = f32[2,2]
- add(%constant.f32.1, %convert.1),
+ %add.1 = f32[]
+ add(%constant.f32, %convert.1),
sharding={maximal device=0}
- %cross-replica-sum.1 = f32[2,2]
+ %cross-replica-sum.1 = f32[]
cross-replica-sum(%add.1),
replica_groups={{0,1}},
to_apply=%sum.f32,
sharding={maximal device=0}
- %cross-replica-sum.ar.2 = bf16[2,2]
+ %cross-replica-sum.ar.2 = bf16[]
cross-replica-sum(%constant.bf16),
replica_groups={{0},{1}},
all_reduce_id=1,
- to_apply=%binary_add,
+ to_apply=%sum.bf16,
sharding={maximal device=1}
- %convert.2 = f32[2,2]
+ %convert.2 = f32[]
convert(%cross-replica-sum.ar.2),
sharding={maximal device=1}
- %add.2 = f32[2,2]
- add(%constant.f32.2, %convert.2),
+ %add.2 = f32[]
+ add(%constant.f32, %convert.2),
sharding={maximal device=1}
- %cross-replica-sum.2 = f32[2,2]
+ %cross-replica-sum.2 = f32[]
cross-replica-sum(%add.2),
replica_groups={{0,1}},
to_apply=%sum.f32,
sharding={maximal device=1}
- ROOT %tuple = (f32[2,2], f32[2,2])
+ ROOT %tuple = (f32[], f32[])
+ tuple(%cross-replica-sum.1, %cross-replica-sum.2),
+ sharding={{maximal device=0}, {maximal device=1}}
+}
+)";
+
+ TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+ ParseAndReturnVerifiedModule(module_str));
+ auto crs_before =
+ module->entry_computation()->root_instruction()->operands()[0];
+ auto replica_groups_before = crs_before->replica_groups();
+ ArCrsCombiner combiner(2);
+ auto changed = combiner.Run(module.get()).ValueOrDie();
+ EXPECT_TRUE(changed);
+ EXPECT_THAT(
+ module->entry_computation()->root_instruction(),
+ op::Tuple(
+ op::CrossReplicaSum(op::Add(
+ op::Divide(op::Constant(), op::Constant()), op::Convert())),
+ op::CrossReplicaSum(op::Add(
+ op::Divide(op::Constant(), op::Constant()), op::Convert()))));
+ auto crs_after =
+ module->entry_computation()->root_instruction()->operands()[0];
+ auto replica_groups_after = crs_after->replica_groups();
+ CompareReplicaGroups(replica_groups_before, replica_groups_after);
+}
+
+TEST_F(ArCrsCombinerTest, OtherSummandNotTheSameDontRewrite) {
+ const char* module_str = R"(
+HloModule foobar
+
+%sum.bf16 (a: bf16[], b: bf16[]) -> bf16[] {
+ %a = bf16[] parameter(0)
+ %b = bf16[] parameter(1)
+ ROOT %add = bf16[] add(%a, %b)
+}
+
+%sum.f32 (x: f32[], y: f32[]) -> f32[] {
+ %x = f32[] parameter(0)
+ %y = f32[] parameter(1)
+ ROOT %add = f32[] add(%x, %y)
+}
+
+ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) {
+ %p = f32[] parameter(0)
+ %constant.bf16 = bf16[] constant(1)
+ %constant.f32.1 = f32[] constant(2)
+ %constant.f32.2 = f32[] constant(3)
+
+ %cross-replica-sum.ar.1 = bf16[]
+ cross-replica-sum(%constant.bf16),
+ replica_groups={{0},{1}},
+ all_reduce_id=1,
+ to_apply=%sum.bf16,
+ sharding={maximal device=0}
+ %convert.1 = f32[]
+ convert(%cross-replica-sum.ar.1),
+ sharding={maximal device=0}
+ %add.1 = f32[]
+ add(%constant.f32.1, %convert.1),
+ sharding={maximal device=0}
+ %cross-replica-sum.1 = f32[]
+ cross-replica-sum(%add.1),
+ replica_groups={{0,1}},
+ to_apply=%sum.f32,
+ sharding={maximal device=0}
+
+ %cross-replica-sum.ar.2 = bf16[]
+ cross-replica-sum(%constant.bf16),
+ replica_groups={{0},{1}},
+ all_reduce_id=1,
+ to_apply=%sum.bf16,
+ sharding={maximal device=1}
+ %convert.2 = f32[]
+ convert(%cross-replica-sum.ar.2),
+ sharding={maximal device=1}
+ %add.2 = f32[]
+ add(%constant.f32.2, %convert.2),
+ sharding={maximal device=1}
+ %cross-replica-sum.2 = f32[]
+ cross-replica-sum(%add.2),
+ replica_groups={{0,1}},
+ to_apply=%sum.f32,
+ sharding={maximal device=1}
+
+ ROOT %tuple = (f32[], f32[])
tuple(%cross-replica-sum.1, %cross-replica-sum.2),
sharding={{maximal device=0}, {maximal device=1}}
}
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index 6374822..f3dfa4d 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -635,18 +635,17 @@
.EmitComputation(
embedded_computation, embedded_computation->name(),
/*is_top_level_computation=*/false,
- &schedule.sequence(embedded_computation).instructions())
+ schedule.sequence(embedded_computation).instructions())
.status());
}
string function_name_prefix = entry_computation->name().empty()
? "__compute"
: entry_computation->name();
- TF_ASSIGN_OR_RETURN(
- llvm::Function * entry_function,
- ir_emitter.EmitComputation(
- entry_computation, function_name_prefix,
- /*is_top_level_computation=*/true,
- &schedule.sequence(entry_computation).instructions()));
+ TF_ASSIGN_OR_RETURN(llvm::Function * entry_function,
+ ir_emitter.EmitComputation(
+ entry_computation, function_name_prefix,
+ /*is_top_level_computation=*/true,
+ schedule.sequence(entry_computation).instructions()));
string function_name = [&]() {
llvm::SmallVector<char, 40> function_name_vector;
@@ -835,7 +834,7 @@
.EmitComputation(
embedded_computation, embedded_computation->name(),
/*is_top_level_computation=*/false,
- &schedule.sequence(embedded_computation).instructions())
+ schedule.sequence(embedded_computation).instructions())
.status());
}
const string& entry_point_name = options.entry_point_name();
@@ -843,7 +842,7 @@
ir_emitter.EmitComputation(
computation, entry_point_name,
/*is_top_level_computation=*/true,
- &schedule.sequence(computation).instructions()));
+ schedule.sequence(computation).instructions()));
CHECK(entry_function->getName() == llvm_ir::AsStringRef(entry_point_name));
diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
index 97f9b85..a33035a 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
@@ -323,11 +323,11 @@
int64 column_remainder = k() % tile_cols();
int64 column_limit = k() - column_remainder;
- ksl_.ForReturnVoid("dot.outer.tiled",
- /*start=*/0, /*end=*/column_limit, /*step=*/tile_cols(),
- [&](llvm::Value* column, bool is_first_column) {
- EmitOuterLoopBody(column, tile_cols(), is_first_column);
- });
+ ksl_.For("dot.outer.tiled",
+ /*start=*/0, /*end=*/column_limit, /*step=*/tile_cols(),
+ [&](llvm::Value* column, bool is_first_column) {
+ EmitOuterLoopBody(column, tile_cols(), is_first_column);
+ });
if (column_remainder != 0) {
EmitOuterLoopBody(b_->getInt64(column_limit), column_remainder,
@@ -340,7 +340,7 @@
int64 columns, bool is_first_column) {
int64 row_limit = m() - (m() % tile_rows());
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.inner.tiled", /*start=*/0, /*end=*/row_limit,
/*step=*/tile_rows(), [&](llvm::Value* row) {
std::vector<llvm::Value*> lhs_tile =
@@ -372,7 +372,7 @@
// // initialized.
// }
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.inner.epilg.outer", /*start=*/current_tile_col,
/*end=*/b_->CreateAdd(columns_llvm, current_tile_col),
/*step=*/1, /*peel_first_iteration=*/false,
@@ -381,14 +381,14 @@
llvm::Value* total_offset = b_->CreateMul(col, b_->getInt64(m()));
llvm::Value* lhs_base_pointer =
vsl_.ComputeOffsetPointer(lhs_, total_offset);
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.inner.epilg.inner", /*start=*/row_start, /*end=*/m(),
/*step=*/1, [&](llvm::Value* scalar_row) {
llvm::Value* product = vsl_.Mul(
vsl_.LoadScalar(lhs_base_pointer, scalar_row), rhs_element);
llvm::Value* setting_result_first_time = b_->CreateAnd(
is_first_scalar_col, b_->getInt1(is_first_tiled_column));
- ksl_.IfReturnVoid(
+ ksl_.If(
setting_result_first_time,
/*true_block_generator=*/
[&]() {
@@ -568,10 +568,9 @@
int64 row_remainder = m() % tile_rows();
int64 row_limit = m() - row_remainder;
- ksl_.ForReturnVoid(
- "dot.outer.tiled",
- /*start=*/0, /*end=*/row_limit, /*step=*/tile_rows(),
- [&](llvm::Value* row) { EmitOuterLoopBody(row, tile_rows()); });
+ ksl_.For("dot.outer.tiled",
+ /*start=*/0, /*end=*/row_limit, /*step=*/tile_rows(),
+ [&](llvm::Value* row) { EmitOuterLoopBody(row, tile_rows()); });
if (row_remainder != 0) {
EmitOuterLoopBody(b_->getInt64(row_limit), row_remainder);
@@ -583,17 +582,17 @@
std::vector<VectorVariable>* vector_accumulators) {
int64 column_limit = k() - (k() % tile_cols());
- ksl_.ForReturnVoid("dot.inner.tiled", /*start=*/0, /*end=*/column_limit,
- /*step=*/tile_cols(), [&](llvm::Value* col) {
- std::vector<llvm::Value*> lhs_tile =
- lhs_memory_tile->LoadTile(/*minor_dim_offset=*/col);
- llvm::Value* rhs_value = vsl_.LoadVector(rhs_, col);
- for (int i = 0; i < rows; i++) {
- llvm::Value* old_sum = (*vector_accumulators)[i].Get();
- (*vector_accumulators)[i].Set(vsl_.Add(
- old_sum, vsl_.Mul(rhs_value, lhs_tile[i])));
- }
- });
+ ksl_.For("dot.inner.tiled", /*start=*/0, /*end=*/column_limit,
+ /*step=*/tile_cols(), [&](llvm::Value* col) {
+ std::vector<llvm::Value*> lhs_tile =
+ lhs_memory_tile->LoadTile(/*minor_dim_offset=*/col);
+ llvm::Value* rhs_value = vsl_.LoadVector(rhs_, col);
+ for (int i = 0; i < rows; i++) {
+ llvm::Value* old_sum = (*vector_accumulators)[i].Get();
+ (*vector_accumulators)[i].Set(
+ vsl_.Add(old_sum, vsl_.Mul(rhs_value, lhs_tile[i])));
+ }
+ });
}
void RowMajorMatrixVectorProductEmitter::EmitInnerLoopEpilogue(
@@ -609,7 +608,7 @@
b_->CreateAdd(b_->getInt64(r), current_tile_row), b_->getInt64(k()));
llvm::Value* lhs_base_pointer =
vsl_.ComputeOffsetPointer(lhs_, total_offset);
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.inner.epilg.inner", /*start=*/column_start, /*end=*/k(),
/*step=*/1, [&](llvm::Value* scalar_col) {
llvm::Value* product =
@@ -813,7 +812,7 @@
if (n_start != dims().n()) {
VectorSupportLibrary vsl(scalar_type(), 1, b_, "gemm");
- ksl_.ForReturnVoid("epi.n", n_start, dims().n(), 1, [&](llvm::Value* n_i) {
+ ksl_.For("epi.n", n_start, dims().n(), 1, [&](llvm::Value* n_i) {
llvm::Value* n_i_next = b_->CreateAdd(n_i, b_->getInt64(1));
HandleResiduesOnK(&vsl, n_i, n_i_next);
});
@@ -924,7 +923,7 @@
VectorSupportLibrary* vsl, int64 tile_size_k, llvm::Value* k_start,
llvm::Value* k_end, llvm::Value* n_start, llvm::Value* n_end,
int64 tile_size_m, llvm::Value* m_start, llvm::Value* m_end) {
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.m", m_start, m_end, tile_size_m, [&](llvm::Value* m_i) {
MemoryTile result_memory_tile(
vsl, b_, /*matrix=*/result_,
@@ -935,11 +934,11 @@
/*matrix_size_along_minor_dim=*/dims().k(),
/*major_dim_offset=*/m_i,
/*tile_size_along_major_dim=*/tile_size_m);
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.n", n_start, n_end, vsl->vector_size(), [&](llvm::Value* n_i) {
TileVariable result_tile_var(vsl,
result_memory_tile.LoadTile(n_i));
- ksl_.ForReturnVoid(
+ ksl_.For(
"dot.k", k_start, k_end, tile_size_k, [&](llvm::Value* k_i) {
MemoryTile rhs_memory_tile(vsl, b_, rhs_, dims().n(), k_i,
tile_size_k);
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index 4032c2d..62a4e8d 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -111,10 +111,9 @@
StatusOr<llvm::Function*> IrEmitter::EmitComputation(
HloComputation* computation, const string& function_name_prefix,
bool is_top_level_computation,
- const std::vector<HloInstruction*>* instruction_order) {
+ absl::Span<HloInstruction* const> instruction_order) {
string function_name = name_uniquer_.GetUniqueName(function_name_prefix);
- VLOG(2) << "Emitting IR for CPU function [" << function_name_prefix
- << "]; ordered? " << (instruction_order != nullptr);
+ VLOG(2) << "Emitting IR for CPU function [" << function_name_prefix << "]";
is_top_level_computation_ = is_top_level_computation;
num_dynamic_loop_bounds_ = 0;
if (!computation->root_instruction()->outer_dimension_partitions().empty()) {
@@ -141,11 +140,7 @@
bool use_rdtscp = arch_type_ == llvm::Triple::ArchType::x86 ||
arch_type_ == llvm::Triple::ArchType::x86_64;
profiling_state_ = ProfilingState(use_rdtscp);
- if (instruction_order == nullptr) {
- TF_RETURN_IF_ERROR(computation->Accept(this));
- } else {
- TF_RETURN_IF_ERROR(computation->AcceptOrdered(this, *instruction_order));
- }
+ TF_RETURN_IF_ERROR(computation->AcceptOrdered(this, instruction_order));
llvm::Function* ir_function = compute_function_->function();
InsertOrDie(&emitted_functions_, computation, ir_function);
// Delete 'compute_function', finalizing 'ir_function' and restoring caller
@@ -2271,6 +2266,22 @@
/*isVarArg=*/false)));
TF_RETURN_IF_ERROR(EmitTargetAddressForOp(custom_call));
+ // Write the tuple table if the output is a tuple.
+ if (ShapeUtil::IsTuple(custom_call->shape())) {
+ std::vector<llvm::Value*> base_ptrs;
+ for (int i = 0; i < ShapeUtil::TupleElementCount(custom_call->shape());
+ ++i) {
+ const Shape& elem_shape =
+ ShapeUtil::GetTupleElementShape(custom_call->shape(), i);
+ TF_RET_CHECK(!ShapeUtil::IsTuple(elem_shape))
+ << "Nested tuples not implemented";
+ TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice slice,
+ assignment_.GetUniqueSlice(custom_call, {i}));
+ llvm::Value* addr = EmitBufferPointer(slice, elem_shape);
+ base_ptrs.push_back(addr);
+ }
+ llvm_ir::EmitTuple(GetIrArrayFor(custom_call), base_ptrs, &b_, module_);
+ }
auto* output_address_arg =
PointerCast(GetEmittedValueFor(custom_call), i8_ptr_type);
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
index 559a816..1db75cc 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
@@ -101,7 +101,7 @@
StatusOr<llvm::Function*> EmitComputation(
HloComputation* computation, const string& function_name_prefix,
bool is_top_level_computation,
- const std::vector<HloInstruction*>* instruction_order);
+ absl::Span<HloInstruction* const> instruction_order);
llvm::IRBuilder<>* b() { return &b_; }
diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc
index efccade..296f39a 100644
--- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc
+++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc
@@ -139,7 +139,7 @@
}
if (func_addr == nullptr) {
- VLOG(2) << "Unable to resolve runtime symbol: " << name;
+ LOG(ERROR) << "Unable to resolve runtime symbol: " << name;
return nullptr;
}
llvm::JITEvaluatedSymbol symbol_info(reinterpret_cast<uint64_t>(func_addr),
@@ -296,6 +296,9 @@
REGISTER_LIBM_SYMBOL(sin, double (*)(double));
#ifdef __APPLE__
REGISTER_LIBM_SYMBOL(__sincos, void (*)(double, double*, double*));
+ registry->Register("__sincosf_stret",
+ reinterpret_cast<void*>(__sincosf_stret));
+ registry->Register("__sincos_stret", reinterpret_cast<void*>(__sincos_stret));
#else
REGISTER_LIBM_SYMBOL(sincos, void (*)(double, double*, double*));
#endif
@@ -311,6 +314,13 @@
registry->Register("memcpy", reinterpret_cast<void*>(memcpy));
registry->Register("memmove", reinterpret_cast<void*>(memmove));
registry->Register("memset", reinterpret_cast<void*>(memset));
+
+#ifdef __APPLE__
+ registry->Register("__bzero", reinterpret_cast<void*>(bzero));
+ registry->Register("memset_pattern16",
+ reinterpret_cast<void*>(memset_pattern16));
+#endif
+
return true;
}
diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc b/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc
new file mode 100644
index 0000000..6d04726
--- /dev/null
+++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc
@@ -0,0 +1,459 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h"
+#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+
+namespace xla {
+
+namespace {
+bool IsTrivialWindowDimension(const WindowDimension& window_dimension) {
+ return window_dimension.size() == 1 && window_dimension.stride() == 1 &&
+ window_dimension.padding_low() == 0 &&
+ window_dimension.padding_high() == 0 &&
+ window_dimension.window_dilation() == 1 &&
+ window_dimension.base_dilation() == 1;
+}
+} // namespace
+
+class DynamicDimensionInferenceVisitor : public DfsHloVisitorWithDefault {
+ public:
+ explicit DynamicDimensionInferenceVisitor(
+ const DynamicParameterBinding& param_bindings,
+ DynamicDimensionInference* parent)
+ : param_bindings_(param_bindings), parent_(parent) {}
+
+ Status DefaultAction(HloInstruction* hlo) override;
+
+ static Status Run(HloComputation* computation,
+ const DynamicParameterBinding& param_bindings,
+ DynamicDimensionInference* parent) {
+ DynamicDimensionInferenceVisitor visitor(param_bindings, parent);
+ return computation->Accept(&visitor);
+ }
+
+ Status HandleParameter(HloInstruction* hlo) override;
+
+ Status HandleReduce(HloInstruction* hlo) override;
+
+ Status HandleDot(HloInstruction* hlo) override;
+
+ Status HandleTranspose(HloInstruction* hlo) override;
+
+ Status HandleReshape(HloInstruction* hlo) override;
+
+ Status HandlePad(HloInstruction* hlo) override;
+
+ Status HandleBroadcast(HloInstruction* hlo) override;
+
+ Status HandleGetDimensionSize(HloInstruction* hlo) override;
+
+ Status HandleSelect(HloInstruction* hlo) override;
+
+ Status HandleConvolution(HloInstruction* hlo) override;
+
+ Status HandleReduceWindow(HloInstruction* hlo) override;
+
+ Status HandleSelectAndScatter(HloInstruction* hlo) override;
+
+ Status HandleGetTupleElement(HloInstruction* hlo) override;
+
+ Status HandleElementwiseUnary(HloInstruction* hlo) override;
+
+ Status HandleElementwiseBinary(HloInstruction* hlo) override;
+
+ private:
+ using OperandDynamicDimensionFn = std::function<Status(
+ HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size)>;
+
+ Status ForEachOperandDynamicDimension(HloInstruction* inst,
+ const OperandDynamicDimensionFn&);
+
+ // Pass through a dynamic dimension from the input to the output with the same
+ // value and index in the shape. This is a helper function to handle trivial
+ // instructions like elementwise operations.
+ Status PassThroughDynamicDimension(HloInstruction*);
+
+ // The dynamic parameter bindings of this computation.
+ const DynamicParameterBinding& param_bindings_;
+
+ // A pointer to DynamicDimensionInference, used to update the dynamic mapping.
+ DynamicDimensionInference* parent_;
+};
+
+Status DynamicDimensionInferenceVisitor::DefaultAction(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ return UnimplementedStrCat(
+ "Asked to propagate a dynamic dimension from hlo ",
+ operand->ToString(), "@", index.ToString(), "@", dimension,
+ " to hlo ", hlo->ToString(), ", which is not implemented.");
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleGetTupleElement(
+ HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ if (hlo->tuple_index() == index[0]) {
+ ShapeIndex new_index =
+ ShapeIndexView(index).ConsumeFront().ToShapeIndex();
+ parent_->SetDynamicSize(hlo, new_index, dimension, dynamic_size);
+ }
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleBroadcast(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ int64 broadcast_dim = hlo->dimensions(dimension);
+ parent_->SetDynamicSize(hlo, index, broadcast_dim, dynamic_size);
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandlePad(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ if (operand_index != 0) {
+ return Unimplemented(
+ "Dynamic dimension on padding value is not supported");
+ }
+ const PaddingConfig_PaddingConfigDimension& padding_config =
+ hlo->padding_config().dimensions(dimension);
+ if (padding_config.interior_padding() == 0 &&
+ padding_config.edge_padding_low() == 0 &&
+ padding_config.edge_padding_high() == 0) {
+ parent_->SetDynamicSize(hlo, {}, dimension, dynamic_size);
+ return Status::OK();
+ } else {
+ return Unimplemented(
+ "Dynamic dimension propagation on padding dimension is not "
+ "supported.");
+ }
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleReduce(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ HloInstruction* reduce = hlo;
+ int64 operand_count = reduce->operand_count();
+ CHECK_EQ(operand_count % 2, 0);
+ if (operand_index >= operand_count / 2) {
+ // Init values doesn't have dynamic size.
+ return Status::OK();
+ }
+ if ((absl::c_count(reduce->dimensions(), dimension) != 0)) {
+ // Dimension is to be reduce, stop tracing.
+ return Status::OK();
+ }
+
+ // Find out the new dynamic dimension after reduce.
+ int64 dimensions_not_reduced_count = 0;
+ for (int i = 0; i < ShapeUtil::Rank(operand->shape()); ++i) {
+ if (dimension == i) {
+ parent_->SetDynamicSize(reduce, {}, dimensions_not_reduced_count,
+ dynamic_size);
+
+ return Status::OK();
+ }
+ if (absl::c_count(reduce->dimensions(), i) == 0) {
+ dimensions_not_reduced_count++;
+ }
+ }
+
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleDot(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ HloInstruction* dot = hlo;
+ const DotDimensionNumbers& dimension_numbers =
+ dot->dot_dimension_numbers();
+ // A map from the operand dimensions to result dimension.
+ absl::flat_hash_map<int64, int64> result_dim_mapping;
+ int64 current_result_dims = 0;
+ std::unordered_set<int64> batch_dims(
+ dimension_numbers.rhs_batch_dimensions().begin(),
+ dimension_numbers.rhs_batch_dimensions().end());
+
+ for (int64 i : dimension_numbers.rhs_batch_dimensions()) {
+ result_dim_mapping[i] = current_result_dims++;
+ }
+
+ for (int64 i = 0; i < ShapeUtil::Rank(dot->operand(0)->shape()); i++) {
+ if (!absl::c_linear_search(
+ dimension_numbers.lhs_contracting_dimensions(), i)) {
+ if (operand_index == 0) {
+ result_dim_mapping[i] = current_result_dims;
+ }
+ current_result_dims++;
+ }
+ }
+
+ for (int64 i = 0; i < ShapeUtil::Rank(dot->operand(1)->shape()); i++) {
+ if (!absl::c_linear_search(
+ dimension_numbers.rhs_contracting_dimensions(), i) &&
+ !absl::c_linear_search(dimension_numbers.rhs_batch_dimensions(),
+ i)) {
+ if (operand_index == 1) {
+ result_dim_mapping[i] = current_result_dims;
+ }
+ current_result_dims++;
+ }
+ }
+
+ // Check if the operand dim is in the result shape. If so, add another
+ // work item to trace that dimension.
+ auto iter = result_dim_mapping.find(dimension);
+ if (iter != result_dim_mapping.end()) {
+ parent_->SetDynamicSize(dot, {}, iter->second, dynamic_size);
+ }
+
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleTranspose(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ parent_->SetDynamicSize(hlo, {}, hlo->dimensions()[dimension],
+ dynamic_size);
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleConvolution(
+ HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ HloInstruction* conv = hlo;
+ const ConvolutionDimensionNumbers& dimension_numbers =
+ conv->convolution_dimension_numbers();
+
+ if (operand_index == 0) {
+ if (dimension == dimension_numbers.input_batch_dimension()) {
+ parent_->SetDynamicSize(conv, {},
+ dimension_numbers.output_batch_dimension(),
+ dynamic_size);
+ return Status::OK();
+ }
+
+ if (dimension == dimension_numbers.input_feature_dimension()) {
+ return Status::OK();
+ }
+ } else {
+ if (dimension == dimension_numbers.kernel_input_feature_dimension()) {
+ return Status::OK();
+ }
+ }
+
+ return Unimplemented("Dynamic Spatial Convolution is not supported: %s",
+ conv->ToString());
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleGetDimensionSize(
+ HloInstruction*) {
+ // Dynamic dimension doesn't propagate through GetDimensionSize:
+ //
+ // Input: F32[x, y, z]
+ // |
+ // GetDimensionSize(1): U32[]
+ //
+ // The returned value is a scalar, which doesn't have any dynamic dimension in
+ // the shape (although the value contains the real size of the dynamic
+ // dimension of the input).
+ return Status::OK();
+}
+
+Status DynamicDimensionInferenceVisitor::PassThroughDynamicDimension(
+ HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ parent_->SetDynamicSize(hlo, index, dimension, dynamic_size);
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleElementwiseUnary(
+ HloInstruction* hlo) {
+ return PassThroughDynamicDimension(hlo);
+}
+
+Status DynamicDimensionInferenceVisitor::HandleSelect(HloInstruction* hlo) {
+ return PassThroughDynamicDimension(hlo);
+}
+
+Status DynamicDimensionInferenceVisitor::HandleElementwiseBinary(
+ HloInstruction* hlo) {
+ return PassThroughDynamicDimension(hlo);
+}
+
+Status DynamicDimensionInferenceVisitor::HandleReshape(HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ HloInstruction* reshape = hlo;
+ std::vector<std::pair<int64, int64>> unmodified_dims =
+ ShapeUtil::DimensionsUnmodifiedByReshape(operand->shape(),
+ reshape->shape());
+ for (auto& unmodified : unmodified_dims) {
+ if (unmodified.first == dimension) {
+ parent_->SetDynamicSize(reshape, {}, unmodified.second,
+ dynamic_size);
+ return Status::OK();
+ }
+ }
+ return Unimplemented(
+ "Dynamic Reshape on modified dimensions is yet not supported: %s",
+ reshape->ToString());
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleReduceWindow(
+ HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ HloInstruction* reduce_window = hlo;
+ const WindowDimension& window_dimension =
+ reduce_window->window().dimensions(dimension);
+
+ if (!IsTrivialWindowDimension(window_dimension)) {
+ return Unimplemented(
+ "Dynamic Spatial reduce window is not supported: %s",
+ reduce_window->ToString());
+ }
+
+ parent_->SetDynamicSize(reduce_window, {}, dimension, dynamic_size);
+
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleSelectAndScatter(
+ HloInstruction* hlo) {
+ return ForEachOperandDynamicDimension(
+ hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension,
+ int64 operand_index, HloInstruction* dynamic_size) {
+ HloInstruction* select_and_scatter = hlo;
+ const WindowDimension& window_dimension =
+ select_and_scatter->window().dimensions(dimension);
+
+ if (!IsTrivialWindowDimension(window_dimension)) {
+ return Unimplemented(
+ "Dynamic Spatial select and scatter is not supported: %s",
+ select_and_scatter->ToString());
+ }
+
+ parent_->SetDynamicSize(select_and_scatter, {}, dimension,
+ dynamic_size);
+
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) {
+ return param_bindings_.ForEachBinding(
+ [&](const DynamicParameterBinding::DynamicParameter& dynamic_parameter,
+ const DynamicParameterBinding::DynamicDimension& dynamic_dimension) {
+ if (dynamic_dimension.parameter_num != hlo->parameter_number()) {
+ return Status::OK();
+ }
+ HloComputation* computation = hlo->parent();
+ HloInstruction* target_parameter =
+ computation->parameter_instruction(dynamic_dimension.parameter_num);
+
+ HloInstruction* dynamic_size =
+ computation->parameter_instruction(dynamic_parameter.parameter_num);
+ for (int64 i : dynamic_parameter.parameter_index) {
+ dynamic_size =
+ computation->AddInstruction(HloInstruction::CreateGetTupleElement(
+ ShapeUtil::GetSubshape(dynamic_size->shape(), {i}),
+ dynamic_size, i));
+ }
+
+ parent_->SetDynamicSize(target_parameter,
+ dynamic_dimension.parameter_index,
+ dynamic_dimension.dimension, dynamic_size);
+ return Status::OK();
+ });
+}
+
+Status DynamicDimensionInferenceVisitor::ForEachOperandDynamicDimension(
+ HloInstruction* inst, const OperandDynamicDimensionFn& fn) {
+ for (int64 operand_index = 0; operand_index < inst->operand_count();
+ ++operand_index) {
+ auto iter =
+ parent_->per_hlo_dynamic_dimensions_.find(inst->operand(operand_index));
+ if (iter != parent_->per_hlo_dynamic_dimensions_.end()) {
+ for (auto& dynamic_dimension : iter->second) {
+ HloInstruction* dynamic_size = parent_->GetDynamicSize(
+ dynamic_dimension.inst, dynamic_dimension.index,
+ dynamic_dimension.dim);
+ TF_RETURN_IF_ERROR(fn(dynamic_dimension.inst, dynamic_dimension.index,
+ dynamic_dimension.dim, operand_index,
+ dynamic_size));
+ }
+ }
+ }
+ return Status::OK();
+}
+
+/* static */
+StatusOr<DynamicDimensionInference> DynamicDimensionInference::Run(
+ HloModule* module) {
+ VLOG(0) << "Param Config " << module->dynamic_parameter_binding().ToString();
+ DynamicDimensionInference inference(module);
+ TF_RETURN_IF_ERROR(inference.AnalyzeDynamicDimensions());
+ return inference;
+}
+
+DynamicDimensionInference::DynamicDimensionInference(HloModule* module)
+ : module_(module) {}
+
+Status DynamicDimensionInference::AnalyzeDynamicDimensions() {
+ return DynamicDimensionInferenceVisitor::Run(
+ module_->entry_computation(), module_->dynamic_parameter_binding(), this);
+}
+
+HloInstruction* DynamicDimensionInference::GetDynamicSize(
+ HloInstruction* inst, const ShapeIndex& index, int64 dim) const {
+ auto iter = dynamic_mapping_.find(DynamicDimension{inst, index, dim});
+ if (iter != dynamic_mapping_.end()) {
+ return iter->second;
+ }
+ return nullptr;
+}
+
+} // namespace xla
diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference.h b/tensorflow/compiler/xla/service/dynamic_dimension_inference.h
new file mode 100644
index 0000000..164d15b
--- /dev/null
+++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference.h
@@ -0,0 +1,112 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_DIMENSION_INFERENCE_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_DIMENSION_INFERENCE_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/status.h"
+#include "tensorflow/compiler/xla/statusor.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/core/platform/macros.h"
+
+namespace xla {
+
+// DynamicDimensionInference analyzes each HLO instruction in a graph and
+// inferences which dimensions are dynamic and which scalar instructions
+// represent the runtime real size of those dynamic dimensions.
+class DynamicDimensionInference {
+ public:
+ static StatusOr<DynamicDimensionInference> Run(HloModule* module);
+
+ string ToString() const;
+
+ // If the dimension `dim` of instruction `inst` at `index` has a dynamic size,
+ // returns a scalar HloInstruction that represents the runtime size of that
+ // dimension. Otherwise returns nullptr.
+ HloInstruction* GetDynamicSize(HloInstruction* inst, const ShapeIndex& index,
+ int64 dim) const;
+
+ friend class DynamicDimensionInferenceVisitor;
+
+ private:
+ explicit DynamicDimensionInference(HloModule* module);
+
+ // DynamicDimension is used as a key in the dynamic key-value mapping. It
+ // unambiguously represents a dynamic dimension of a instruction at a given
+ // index.
+ struct DynamicDimension {
+ // HloInstruction that holds the dimension.
+ HloInstruction* inst;
+ // Subshape of the instruction that holds the dimension.
+ ShapeIndex index;
+ // The dimension number of the dynamic dimension at given index of a given
+ // instruction.
+ int64 dim;
+
+ // Artifacts needed to make this struct able to be used as a `key` in absl
+ // maps. "friend" keywords are added so these functions can be found through
+ // ADL.
+ template <typename H>
+ friend H AbslHashValue(H h, const DynamicDimension& m) {
+ return H::combine(std::move(h), m.inst, m.index, m.dim);
+ }
+
+ friend bool operator==(const DynamicDimension& lhs,
+ const DynamicDimension& rhs) {
+ return lhs.inst == rhs.inst && lhs.index == rhs.index &&
+ lhs.dim == rhs.dim;
+ }
+ };
+
+ // Update the dynamic mapping so that we know dimension `dim` of instruction
+ // `inst` at `index` has a dynamic size, and its runtime size is represented
+ // by a scalar instruction `size`.
+ void SetDynamicSize(HloInstruction* inst, const ShapeIndex& index, int64 dim,
+ HloInstruction* size) {
+ dynamic_mapping_.try_emplace(DynamicDimension{inst, index, dim}, size);
+ auto iter = per_hlo_dynamic_dimensions_.try_emplace(inst);
+ iter.first->second.emplace(DynamicDimension{inst, index, dim});
+ }
+
+ // AnalyzeDynamicDimensions starts the analysis of the dynamic dimensions in
+ // module_.
+ Status AnalyzeDynamicDimensions();
+
+ // HloModule being analyzed.
+ HloModule* module_;
+
+ // dynamic_mapping_ holds the result of the analysis. It maps a dynamic
+ // dimension to a scalar HloInstruction that represents the real dynamic size
+ // of the dynamic dimension.
+ using DynamicMapping = absl::flat_hash_map<DynamicDimension, HloInstruction*>;
+ DynamicMapping dynamic_mapping_;
+
+ using PerHloDynamicDimensions =
+ absl::flat_hash_map<HloInstruction*,
+ absl::flat_hash_set<DynamicDimension>>;
+ PerHloDynamicDimensions per_hlo_dynamic_dimensions_;
+};
+
+} // namespace xla
+
+#endif // TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_DIMENSION_INFERENCE_H_
diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc b/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc
new file mode 100644
index 0000000..ea9ebed
--- /dev/null
+++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc
@@ -0,0 +1,535 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h"
+
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/compiler/xla/literal.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_matchers.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/hlo_opcode.h"
+#include "tensorflow/compiler/xla/service/hlo_runner.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"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace op = xla::testing::opcode_matchers;
+
+namespace xla {
+namespace {
+
+class DynamicDimensionInferenceTest : public HloTestBase {
+ protected:
+ DynamicDimensionInferenceTest() : HloTestBase() {
+ module_ = CreateNewVerifiedModule();
+ }
+
+ Status RunInference() {
+ hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before alias analysis");
+ TF_ASSIGN_OR_RETURN(DynamicDimensionInference inference,
+ DynamicDimensionInference::Run(module_.get()));
+
+ inference_ = absl::make_unique<DynamicDimensionInference>(inference);
+ return Status::OK();
+ }
+
+ HloComputation* GetAdd() {
+ auto embedded_builder = HloComputation::Builder("add");
+ auto lhs = embedded_builder.AddInstruction(HloInstruction::CreateParameter(
+ 0, ShapeUtil::MakeShape(F32, {}), "lhs"));
+ auto rhs = embedded_builder.AddInstruction(HloInstruction::CreateParameter(
+ 1, ShapeUtil::MakeShape(F32, {}), "rhs"));
+ embedded_builder.AddInstruction(
+ HloInstruction::CreateBinary(lhs->shape(), HloOpcode::kAdd, lhs, rhs));
+ return module_->AddEmbeddedComputation(embedded_builder.Build());
+ }
+
+ std::unique_ptr<HloModule> module_;
+ std::unique_ptr<DynamicDimensionInference> inference_;
+ const Shape scalar_shape_ = ShapeUtil::MakeShape(S32, {});
+};
+
+TEST_F(DynamicDimensionInferenceTest, ParamTest) {
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2});
+
+ auto param = builder.AddInstruction(
+ HloInstruction::CreateParameter(0, input_shape, "param"));
+ auto param2 = builder.AddInstruction(
+ HloInstruction::CreateParameter(1, scalar_shape_, "param"));
+
+ module_->AddEntryComputation(builder.Build());
+ // Set up dynamic parameter binding.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(param, {}, 1), param2);
+ EXPECT_EQ(inference_->GetDynamicSize(param, {}, 0), nullptr);
+ EXPECT_EQ(inference_->GetDynamicSize(param2, {}, 0), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ParamTestTuple) {
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2});
+
+ auto param = builder.AddInstruction(HloInstruction::CreateParameter(
+ 0, ShapeUtil::MakeTupleShape({input_shape, scalar_shape_}), "param"));
+
+ module_->AddEntryComputation(builder.Build());
+ // Set up dynamic parameter binding.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{0, {1}},
+ DynamicParameterBinding::DynamicDimension{0, {0}, 1}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_THAT(inference_->GetDynamicSize(param, {0}, 1),
+ op::GetTupleElement(param, 1));
+
+ EXPECT_EQ(inference_->GetDynamicSize(param, {0}, 0), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, GetTupleElement) {
+ // When data flows through GTE, the dynamic dimension size keeps the
+ // same, and the index has its front popped.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2});
+
+ auto param = builder.AddInstruction(HloInstruction::CreateParameter(
+ 0, ShapeUtil::MakeTupleShape({input_shape, scalar_shape_}), "param"));
+
+ auto gte = builder.AddInstruction(
+ HloInstruction::CreateGetTupleElement(input_shape, param, 0));
+
+ module_->AddEntryComputation(builder.Build());
+ // Set up dynamic parameter binding.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{0, {1}},
+ DynamicParameterBinding::DynamicDimension{0, {0}, 1}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_THAT(inference_->GetDynamicSize(param, {0}, 1),
+ op::GetTupleElement(param, 1));
+
+ EXPECT_THAT(inference_->GetDynamicSize(gte, {}, 1),
+ op::GetTupleElement(param, 1));
+
+ EXPECT_EQ(inference_->GetDynamicSize(param, {0}, 0), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ElementwiseTest) {
+ // When data flows through elementwise, the dynamic dimension size keeps the
+ // same.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2});
+
+ auto data_param = builder.AddInstruction(
+ HloInstruction::CreateParameter(0, input_shape, "data_param"));
+ auto size_param = builder.AddInstruction(
+ HloInstruction::CreateParameter(1, scalar_shape_, "size_param"));
+
+ auto* negate = builder.AddInstruction(
+ HloInstruction::CreateUnary(input_shape, HloOpcode::kNegate, data_param));
+
+ module_->AddEntryComputation(builder.Build());
+ // Set up dynamic parameter binding.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(negate, {}, 1), size_param);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ReduceTestI) {
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2});
+ auto reduce_shape = ShapeUtil::MakeShape(F32, {2});
+
+ auto data_param = builder.AddInstruction(
+ HloInstruction::CreateParameter(0, input_shape, "data_param"));
+ auto size_param = builder.AddInstruction(
+ HloInstruction::CreateParameter(1, scalar_shape_, "size_param"));
+
+ auto negate = builder.AddInstruction(
+ HloInstruction::CreateUnary(input_shape, HloOpcode::kNegate, data_param));
+
+ auto init = builder.AddInstruction(
+ HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(0.0)));
+
+ auto reduce = builder.AddInstruction(HloInstruction::CreateReduce(
+ reduce_shape, negate, init, {0, 2}, GetAdd()));
+
+ module_->AddEntryComputation(builder.Build());
+
+ // Set up dynamic parameter binding.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(reduce, {}, 0), size_param);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ReduceTestII) {
+ // Same as ReduceTestI, but only reduce one dimension.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2});
+ auto reduce_shape = ShapeUtil::MakeShape(F32, {1, 2});
+
+ auto data_param = builder.AddInstruction(
+ HloInstruction::CreateParameter(0, input_shape, "data_param"));
+ auto size_param = builder.AddInstruction(
+ HloInstruction::CreateParameter(1, scalar_shape_, "size_param"));
+
+ auto negate = builder.AddInstruction(
+ HloInstruction::CreateUnary(input_shape, HloOpcode::kNegate, data_param));
+
+ auto init = builder.AddInstruction(
+ HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(0.0)));
+
+ auto reduce = builder.AddInstruction(
+ HloInstruction::CreateReduce(reduce_shape, negate, init, {1}, GetAdd()));
+
+ module_->AddEntryComputation(builder.Build());
+
+ // Set up dynamic parameter binding.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 2}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(reduce, {}, 1), size_param);
+ EXPECT_EQ(inference_->GetDynamicSize(reduce, {}, 0), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, DotTest) {
+ auto builder = HloComputation::Builder(TestName());
+ constexpr int xdim = 3;
+ constexpr int ydim = 2;
+ constexpr int zdim = 1;
+ auto xy_shape = ShapeUtil::MakeShape(F32, {xdim, ydim});
+ auto yz_shape = ShapeUtil::MakeShape(F32, {ydim, zdim});
+ auto xz_shape = ShapeUtil::MakeShape(F32, {xdim, zdim});
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, xy_shape, "A"));
+ auto* b_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, yz_shape, "B"));
+ auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/2, scalar_shape_, "size_param"));
+
+ DotDimensionNumbers dot_dnums;
+ dot_dnums.add_lhs_contracting_dimensions(1);
+ dot_dnums.add_rhs_contracting_dimensions(0);
+ auto dot = builder.AddInstruction(
+ HloInstruction::CreateDot(xz_shape, a_param, b_param, dot_dnums,
+ HloTestBase::DefaultPrecisionConfig(2)));
+
+ module_->AddEntryComputation(builder.Build());
+
+ // Set up dynamic parameter binding for non-contracting dimension.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{2, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 0}));
+
+ // Set up binding for contracting dimensions.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{2, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{2, {}},
+ DynamicParameterBinding::DynamicDimension{1, {}, 0}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(dot, {}, 0), size_param);
+ EXPECT_EQ(inference_->GetDynamicSize(dot, {}, 1), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ConvolutionTest) {
+ auto builder = HloComputation::Builder(TestName());
+ constexpr int xdim = 3;
+ constexpr int ydim = 2;
+ constexpr int zdim = 1;
+ auto xy_shape = ShapeUtil::MakeShape(F32, {xdim, ydim});
+ auto yz_shape = ShapeUtil::MakeShape(F32, {ydim, zdim});
+ auto zx_shape = ShapeUtil::MakeShape(F32, {zdim, xdim});
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, xy_shape, "A"));
+ auto* b_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, yz_shape, "B"));
+ auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/2, scalar_shape_, "size_param"));
+
+ auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers(0);
+
+ dnums.set_kernel_input_feature_dimension(0);
+ dnums.set_kernel_output_feature_dimension(1);
+ dnums.set_input_batch_dimension(0);
+ dnums.set_output_batch_dimension(1);
+ dnums.set_output_feature_dimension(0);
+
+ Window window;
+
+ auto* conv = builder.AddInstruction(HloInstruction::CreateConvolve(
+ zx_shape, a_param, b_param, /*feature_group_count=*/1, window, dnums,
+ HloTestBase::DefaultPrecisionConfig(2)));
+
+ module_->AddEntryComputation(builder.Build());
+
+ // Set up dynamic parameter binding for non-contracting dimension.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{2, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 0}));
+
+ // Set up binding for contracting dimensions.
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{2, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(conv, {}, 1), size_param);
+ EXPECT_EQ(inference_->GetDynamicSize(conv, {}, 0), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, TransposeTest) {
+ // Test the ability to trace unmodified dimensions
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 3});
+ auto output_shape = ShapeUtil::MakeShape(F32, {3, 2, 1});
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, input_shape, "A"));
+ auto* size_param_1 = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, scalar_shape_, "size_param"));
+ auto* size_param_2 = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/2, scalar_shape_, "size_param"));
+ auto* size_param_3 = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/3, scalar_shape_, "size_param"));
+
+ auto* transpose = builder.AddInstruction(
+ HloInstruction::CreateTranspose(output_shape, a_param, {2, 1, 0}));
+
+ module_->AddEntryComputation(builder.Build());
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 0}));
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{2, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{3, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 2}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(transpose, {}, 0), size_param_3);
+ EXPECT_EQ(inference_->GetDynamicSize(transpose, {}, 1), size_param_2);
+ EXPECT_EQ(inference_->GetDynamicSize(transpose, {}, 2), size_param_1);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ReshapeTest) {
+ // Test the ability to trace unmodified reshape dimensions.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {2, 3, 4, 5, 6});
+ auto output_shape = ShapeUtil::MakeShape(F32, {6, 4, 1, 5, 2, 3});
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, input_shape, "A"));
+ auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, scalar_shape_, "size_param"));
+
+ auto* reshape = builder.AddInstruction(
+ HloInstruction::CreateReshape(output_shape, a_param));
+
+ module_->AddEntryComputation(builder.Build());
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 2}));
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 3}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 0), nullptr);
+ EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 1), size_param);
+ EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 2), nullptr);
+ EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 3), size_param);
+ EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 4), nullptr);
+ EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 5), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ReshapeTestUnimplemented) {
+ // Test the ability to trace unmodified reshape dimensions.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {2, 3, 4, 5, 6});
+ auto output_shape = ShapeUtil::MakeShape(F32, {6, 4, 1, 5, 2, 3});
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, input_shape, "A"));
+
+ builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, scalar_shape_, "size_param"));
+
+ builder.AddInstruction(HloInstruction::CreateReshape(output_shape, a_param));
+
+ module_->AddEntryComputation(builder.Build());
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 1}));
+
+ Status status = RunInference();
+ EXPECT_EQ(status.code(), tensorflow::error::UNIMPLEMENTED);
+}
+
+TEST_F(DynamicDimensionInferenceTest, BroadcastTest) {
+ // Test the ability to trace broadcast dimension.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {2});
+ auto output_shape = ShapeUtil::MakeShape(F32, {3, 2, 4});
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, input_shape, "A"));
+ auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, scalar_shape_, "size_param"));
+
+ auto* broadcast = builder.AddInstruction(
+ HloInstruction::CreateBroadcast(output_shape, a_param, {1}));
+
+ module_->AddEntryComputation(builder.Build());
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 0}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 0), nullptr);
+ EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 1), size_param);
+ EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 2), nullptr);
+}
+
+TEST_F(DynamicDimensionInferenceTest, ReduceWindowBatchTest) {
+ // Test the ability to trace reduce window batch dimensions.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {2, 4, 4});
+ auto output_shape = ShapeUtil::MakeShape(F32, {2, 2, 2});
+
+ Window window;
+ // First dimension is unchanged.
+ WindowDimension* batch_dim = window.add_dimensions();
+ batch_dim->set_size(1);
+ batch_dim->set_stride(1);
+ batch_dim->set_padding_low(0);
+ batch_dim->set_padding_high(0);
+ batch_dim->set_window_dilation(1);
+ batch_dim->set_base_dilation(1);
+
+ // Second and third dimension are reduced.
+ for (int64 i = 0; i < 2; ++i) {
+ WindowDimension* dim = window.add_dimensions();
+ dim->set_size(2);
+ dim->set_stride(2);
+ dim->set_padding_low(0);
+ dim->set_padding_high(0);
+ dim->set_window_dilation(1);
+ dim->set_base_dilation(1);
+ }
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, input_shape, "A"));
+ auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, scalar_shape_, "size_param"));
+
+ auto init = builder.AddInstruction(
+ HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(0.0)));
+
+ auto* reduce_window =
+ builder.AddInstruction(HloInstruction::CreateReduceWindow(
+ output_shape, a_param, init, window, GetAdd()));
+
+ module_->AddEntryComputation(builder.Build());
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 0}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(reduce_window, {}, 0), size_param);
+}
+
+TEST_F(DynamicDimensionInferenceTest, SelectAndScatterTest) {
+ // Test the ability to trace select and scatter batch dimensions.
+ auto builder = HloComputation::Builder(TestName());
+ auto input_shape = ShapeUtil::MakeShape(F32, {2, 4, 4});
+ auto output_shape = ShapeUtil::MakeShape(F32, {2, 2, 2});
+
+ Window window;
+ // First dimension is unchanged.
+ WindowDimension* batch_dim = window.add_dimensions();
+ batch_dim->set_size(1);
+ batch_dim->set_stride(1);
+ batch_dim->set_padding_low(0);
+ batch_dim->set_padding_high(0);
+ batch_dim->set_window_dilation(1);
+ batch_dim->set_base_dilation(1);
+
+ // Second and third dimension are reduced.
+ for (int64 i = 0; i < 2; ++i) {
+ WindowDimension* dim = window.add_dimensions();
+ dim->set_size(2);
+ dim->set_stride(2);
+ dim->set_padding_low(0);
+ dim->set_padding_high(0);
+ dim->set_window_dilation(1);
+ dim->set_base_dilation(1);
+ }
+
+ auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/0, input_shape, "A"));
+ auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter(
+ /*parameter_number=*/1, scalar_shape_, "size_param"));
+
+ auto init = builder.AddInstruction(
+ HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(0.0)));
+
+ auto* reduce_window =
+ builder.AddInstruction(HloInstruction::CreateReduceWindow(
+ output_shape, a_param, init, window, GetAdd()));
+
+ module_->AddEntryComputation(builder.Build());
+
+ TF_CHECK_OK(module_->dynamic_parameter_binding().Bind(
+ DynamicParameterBinding::DynamicParameter{1, {}},
+ DynamicParameterBinding::DynamicDimension{0, {}, 0}));
+
+ TF_ASSERT_OK(RunInference());
+ EXPECT_EQ(inference_->GetDynamicSize(reduce_window, {}, 0), size_param);
+}
+
+} // namespace
+} // namespace xla
diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc
index d782904..3a09d4d 100644
--- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc
+++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc
@@ -43,13 +43,14 @@
// dilation), returns kPad and/or kSlice instructions that explicitly apply the
// padding; otherwise returns the original input operand. When there is both
// positive padding (including dilation) and negative padding, we insert both
-// kPad and kSlice.
+// kPad and kSlice. Modifies 'conv_window' accordingly if any padding was moved
+// into a kPad or kSlice op.
HloInstruction* MaybePaddedAndSlicedInput(
- const Window& conv_window, const ConvolutionDimensionNumbers& conv_dnums,
+ Window* conv_window, const ConvolutionDimensionNumbers& conv_dnums,
HloInstruction* input) {
HloComputation* computation = input->parent();
- if (!window_util::HasSymmetricPadding(conv_window) ||
- window_util::HasBaseDilation(conv_window)) {
+ if (!window_util::HasSymmetricPadding(*conv_window) ||
+ window_util::HasBaseDilation(*conv_window)) {
// If padding is uneven or has dilation, we insert a kPad instruction that
// applies positive padding and dilation.
//
@@ -62,12 +63,21 @@
MakeNoPaddingConfig(input->shape().dimensions_size());
for (size_t i = 0; i < conv_dnums.input_spatial_dimensions().size(); ++i) {
int64 dim = conv_dnums.input_spatial_dimensions(i);
- padding_config.mutable_dimensions(dim)->set_edge_padding_low(
- std::max<int64>(0LL, conv_window.dimensions(i).padding_low()));
- padding_config.mutable_dimensions(dim)->set_edge_padding_high(
- std::max<int64>(0LL, conv_window.dimensions(i).padding_high()));
- padding_config.mutable_dimensions(dim)->set_interior_padding(
- conv_window.dimensions(i).base_dilation() - 1);
+ if (conv_window->dimensions(i).padding_low() > 0) {
+ padding_config.mutable_dimensions(dim)->set_edge_padding_low(
+ conv_window->dimensions(i).padding_low());
+ conv_window->mutable_dimensions(i)->set_padding_low(0);
+ }
+ if (conv_window->dimensions(i).padding_high() > 0) {
+ padding_config.mutable_dimensions(dim)->set_edge_padding_high(
+ conv_window->dimensions(i).padding_high());
+ conv_window->mutable_dimensions(i)->set_padding_high(0);
+ }
+ if (conv_window->dimensions(i).base_dilation() != 1) {
+ padding_config.mutable_dimensions(dim)->set_interior_padding(
+ conv_window->dimensions(i).base_dilation() - 1);
+ conv_window->mutable_dimensions(i)->set_base_dilation(1);
+ }
}
PrimitiveType element_type = input->shape().element_type();
HloInstruction* padding = computation->AddInstruction(
@@ -75,7 +85,7 @@
input = MakePadHlo(input, padding, padding_config).ValueOrDie();
}
- if (window_util::HasNegativePadding(conv_window)) {
+ if (window_util::HasNegativePadding(*conv_window)) {
// If the window has negative padding, insert a kSlice that explicitly
// applies negative padding.
//
@@ -89,10 +99,14 @@
int64 dim = conv_dnums.input_spatial_dimensions(i);
// If dimension "dim" has negative padding, increase the start index or
// decrement the limit index by the amount of negative padding.
- start_indices[dim] +=
- std::max<int64>(0LL, -conv_window.dimensions(i).padding_low());
- limit_indices[dim] -=
- std::max<int64>(0LL, -conv_window.dimensions(i).padding_high());
+ if (conv_window->dimensions(i).padding_low() < 0) {
+ start_indices[dim] += -conv_window->dimensions(i).padding_low();
+ conv_window->mutable_dimensions(i)->set_padding_low(0);
+ }
+ if (conv_window->dimensions(i).padding_high() < 0) {
+ limit_indices[dim] -= -conv_window->dimensions(i).padding_high();
+ conv_window->mutable_dimensions(i)->set_padding_high(0);
+ }
}
input =
@@ -140,25 +154,22 @@
// Insert slices and/or pads between the convolution and its input and/or
// kernel operand.
+ Window new_conv_window = conv->window();
HloInstruction* new_input = MaybePaddedAndSlicedInput(
- conv->window(), conv->convolution_dimension_numbers(),
+ &new_conv_window, conv->convolution_dimension_numbers(),
conv->mutable_operand(0));
HloInstruction* new_kernel =
- MaybePaddedKernel(conv->window(), conv->convolution_dimension_numbers(),
+ MaybePaddedKernel(new_conv_window, conv->convolution_dimension_numbers(),
conv->mutable_operand(1));
- // Remove the padding from convolution's window field. These paddings are
- // made explicit with the inserted pads.
- Window new_conv_window = conv->window();
+ // Remove the window dilation from convolution's window field. These paddings
+ // are made explicit with the pads inserted by MaybePaddedKernel().
for (size_t i = 0; i < new_conv_window.dimensions_size(); ++i) {
WindowDimension* dim = new_conv_window.mutable_dimensions(i);
// The size of the kernel may have changed so update the Window to match.
dim->set_size(new_kernel->shape().dimensions(
conv->convolution_dimension_numbers().kernel_spatial_dimensions(i)));
- dim->set_padding_low(0);
- dim->set_padding_high(0);
- dim->set_base_dilation(1);
dim->set_window_dilation(1);
}
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
index 38b8923..87d16c0 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
@@ -2064,11 +2064,9 @@
// pressure, since we touch threadIdx.x and blockIdx.x at the beginning of the
// kernel *anyway*.
std::vector<IrArray> output_arrays = ConstructIrArrayForOutputs(hlo);
- TF_RETURN_IF_ERROR(
- KernelSupportLibrary(&b_).If("emit_mof_tuple", IsBlock0Thread0(&b_), [&] {
- llvm_ir::EmitTuple(GetIrArray(hlo, hlo), output_arrays, &b_, module_);
- return Status::OK();
- }));
+ KernelSupportLibrary{&b_}.If("emit_mof_tuple", IsBlock0Thread0(&b_), [&] {
+ llvm_ir::EmitTuple(GetIrArray(hlo, hlo), output_arrays, &b_, module_);
+ });
// For multioutput fusion, we need to emit each operand and the root.
TF_RETURN_IF_ERROR(
@@ -2138,34 +2136,36 @@
namespace {
-void EmitFullTile(const KernelMappingScheme* mapping_scheme,
- const IrArray::Index& tile_origin_index,
- llvm::IRBuilder<>* builder, llvm::Value* y, llvm::Value* x,
- llvm::Type* index_ty,
- const std::function<void(const IrArray::Index&, llvm::Value*,
- llvm::Value*)>& emit_elem_function) {
+void EmitFullElementalTile(
+ const KernelMappingScheme* mapping_scheme,
+ const IrArray::Index& tile_origin_index, const string& loop_name,
+ KernelSupportLibrary* ksl, llvm::IRBuilder<>* builder, llvm::Value* y,
+ llvm::Value* x, llvm::Type* index_ty,
+ const std::function<void(const IrArray::Index&, llvm::Value*,
+ llvm::Value*)>& emit_elem_function) {
int64 num_threads_x = mapping_scheme->GetNumberOfThreadsForDimensionX();
int64 num_threads_y = mapping_scheme->GetNumberOfThreadsForDimensionY();
int64 tile_size_x = mapping_scheme->GetTileSizeForDimensionX();
int64 tile_size_y = mapping_scheme->GetTileSizeForDimensionY();
- for (int64 i = 0; i < tile_size_y; i += num_threads_y) {
- IrArray::Index source_idx_y =
- tile_origin_index.AddOffsetToDim(llvm::ConstantInt::get(index_ty, i),
- KernelMappingScheme::DimY, builder);
- llvm::Value* y_loc =
- builder->CreateAdd(llvm::ConstantInt::get(index_ty, i), y);
- for (int64 j = 0; j < tile_size_x; j += num_threads_x) {
- IrArray::Index source_idx =
- source_idx_y.AddOffsetToDim(llvm::ConstantInt::get(index_ty, j),
- KernelMappingScheme::DimX, builder);
- llvm::Value* x_loc =
- builder->CreateAdd(llvm::ConstantInt::get(index_ty, j), x);
- emit_elem_function(source_idx, y_loc, x_loc);
- }
- }
+ ksl->For(loop_name + "_y", /*start=*/llvm::ConstantInt::get(index_ty, 0),
+ /*end=*/llvm::ConstantInt::get(index_ty, tile_size_y),
+ /*step=*/llvm::ConstantInt::get(index_ty, num_threads_y),
+ [&](llvm::Value* y_indvar) {
+ IrArray::Index source_idx_y = tile_origin_index.AddOffsetToDim(
+ y_indvar, KernelMappingScheme::DimY, builder);
+ llvm::Value* y_loc = builder->CreateAdd(y_indvar, y);
+ for (int64 j = 0; j < tile_size_x; j += num_threads_x) {
+ IrArray::Index source_idx = source_idx_y.AddOffsetToDim(
+ llvm::ConstantInt::get(index_ty, j),
+ KernelMappingScheme::DimX, builder);
+ llvm::Value* x_loc =
+ builder->CreateAdd(llvm::ConstantInt::get(index_ty, j), x);
+ emit_elem_function(source_idx, y_loc, x_loc);
+ }
+ });
}
-void EmitPartialTile(
+void EmitPartialElementalTile(
const KernelMappingScheme* mapping_scheme,
const IrArray::Index& tile_origin_index, const string& loop_name,
KernelSupportLibrary* ksl, llvm::IRBuilder<>* builder, llvm::Value* y,
@@ -2184,7 +2184,7 @@
llvm::Value* x_loc =
builder->CreateAdd(llvm::ConstantInt::get(index_ty, j), x);
- ksl->IfReturnVoid(
+ ksl->If(
loop_name + "_x_in_tile", builder->CreateICmpULT(x_loc, tile_width),
[&] {
// tile_height_bound =
@@ -2196,13 +2196,13 @@
llvm::Value* tile_height_bound = builder->CreateMul(
ceiling_of_ratio,
llvm::ConstantInt::get(index_ty, num_threads_y));
- ksl->ForReturnVoid(
+ ksl->For(
loop_name, /*start=*/llvm::ConstantInt::get(index_ty, 0),
/*end=*/tile_height_bound,
/*step=*/llvm::ConstantInt::get(index_ty, num_threads_y),
[&](llvm::Value* y_indvar) {
llvm::Value* y_loc = builder->CreateAdd(y_indvar, y);
- ksl->IfReturnVoid(
+ ksl->If(
loop_name + "_y_in_tile",
builder->CreateICmpULT(y_loc, tile_height), [&] {
emit_elem_function(
@@ -2234,7 +2234,7 @@
int64 tile_size_y = mapping_scheme->GetTileSizeForDimensionY();
llvm::Type* index_ty = tile_width->getType();
- ksl->IfReturnVoid(
+ ksl->If(
loop_name + "_full_tile",
builder->CreateAnd(
builder->CreateICmpEQ(llvm::ConstantInt::get(index_ty, tile_size_x),
@@ -2242,13 +2242,13 @@
builder->CreateICmpEQ(llvm::ConstantInt::get(index_ty, tile_size_y),
tile_height)),
[&] {
- EmitFullTile(mapping_scheme, tile_origin_index, builder, y, x, index_ty,
- emit_elem_function);
+ EmitFullElementalTile(mapping_scheme, tile_origin_index, loop_name, ksl,
+ builder, y, x, index_ty, emit_elem_function);
},
[&] {
- EmitPartialTile(mapping_scheme, tile_origin_index, loop_name, ksl,
- builder, y, x, tile_height, tile_width, index_ty,
- emit_elem_function);
+ EmitPartialElementalTile(mapping_scheme, tile_origin_index, loop_name,
+ ksl, builder, y, x, tile_height, tile_width,
+ index_ty, emit_elem_function);
});
}
} // namespace
@@ -2358,14 +2358,14 @@
AddressVector* GetMutablePartialResultAddresses() {
return &partial_result_addresses_;
}
- const AddressVector& GetPartialResultAddresses() const {
+ absl::Span<llvm::AllocaInst* const> GetPartialResultAddresses() const {
return partial_result_addresses_;
}
AddressVector* GetMutableReductionInputAddresses() {
return &reduction_input_addresses_;
}
- const AddressVector& GetReductionInputAddresses() const {
+ absl::Span<llvm::AllocaInst* const> GetReductionInputAddresses() const {
return reduction_input_addresses_;
}
@@ -2378,7 +2378,7 @@
InlinedVector<ShapeIndex, 1>* GetMutableReductionOutputShapeIndices() {
return &reduction_output_shape_indices_;
}
- const InlinedVector<ShapeIndex, 1>& GetReductionOutputShapeIndices() const {
+ absl::Span<const ShapeIndex> GetReductionOutputShapeIndices() const {
return reduction_output_shape_indices_;
}
@@ -2533,8 +2533,8 @@
}
void IrEmitterUnnested::EmitFullWarpShuffleDownLoopForAllReduces(
- const InlinedVector<HloComputation*, 1>& reducers,
- const AddressVector& partial_result_addresses) {
+ absl::Span<HloComputation* const> reducers,
+ absl::Span<llvm::AllocaInst* const> partial_result_addresses) {
for (int distance = 16; distance >= 1; distance /= 2) {
for (int i = 0; i != reducers.size(); ++i) {
llvm::Type* element_type =
@@ -2566,11 +2566,11 @@
ReductionCodegenInfo* reduction_info =
static_cast<ReductionCodegenInfo*>(kernel_info);
int num_reduces = reduction_info->GetNumberOfReduces();
- const AddressVector& partial_result_addresses =
+ absl::Span<llvm::AllocaInst* const> partial_result_addresses =
reduction_info->GetPartialResultAddresses();
const InlinedVector<HloComputation*, 1>& reducers =
reduction_info->GetReducers();
- const InlinedVector<ShapeIndex, 1>& reduction_output_shape_indices =
+ absl::Span<const ShapeIndex> reduction_output_shape_indices =
reduction_info->GetReductionOutputShapeIndices();
if (reduction_info->IsRowReduction()) {
@@ -2690,9 +2690,9 @@
reduction_info->GetKernelMappingScheme()->GetUnnormalizedIndex(
index,
GetFirstReduceInstruction(output_instructions)->operand(0)->shape());
- const AddressVector& partial_reduction_result_addresses =
+ absl::Span<llvm::AllocaInst* const> partial_reduction_result_addresses =
reduction_info->GetPartialResultAddresses();
- const AddressVector& reduction_input_addresses =
+ absl::Span<llvm::AllocaInst* const> reduction_input_addresses =
reduction_info->GetReductionInputAddresses();
const InlinedVector<HloComputation*, 1>& reducers =
reduction_info->GetReducers();
@@ -2751,15 +2751,14 @@
Select(ICmpEQ(last_block_for_dim, block_id_for_dim),
last_block_size_for_dim, block_size_for_dim);
- ksl.ForReturnVoid(
- loop_name,
- /*start=*/index_typed_constant(0),
- /*end=*/num_tiles_in_block,
- /*step=*/1, [&](llvm::Value* block_dim_induction_var) {
- IrArray::Index tile_index = starting_tile.AddOffsetToDim(
- block_dim_induction_var, dim_id, &b_);
- emit_next_block_dim(tile_index);
- });
+ ksl.For(loop_name,
+ /*start=*/index_typed_constant(0),
+ /*end=*/num_tiles_in_block,
+ /*step=*/1, [&](llvm::Value* block_dim_induction_var) {
+ IrArray::Index tile_index = starting_tile.AddOffsetToDim(
+ block_dim_induction_var, dim_id, &b_);
+ emit_next_block_dim(tile_index);
+ });
}
};
@@ -2841,10 +2840,22 @@
<< llvm_ir::DumpToString(*param_shmem_buffers[id]);
}
- LaunchDimensions launch_dimensions = LaunchDimensions(
- mapping_scheme->GetNumberOfBlocks(), mapping_scheme->GetThreadsPerTile());
- llvm::Type* index_ty = GetIndexTypeForKernel(
- unnested_hlo, launch_dimensions.launch_bound(), &b_);
+ const ReductionCodegenInfo* reduction_info =
+ dynamic_cast<const ReductionCodegenInfo*>(kernel_info);
+ bool is_column_reduction =
+ (reduction_info && !reduction_info->IsRowReduction());
+
+ LaunchDimensions launch_dimensions =
+ LaunchDimensions(mapping_scheme->GetNumberOfBlocks(),
+ mapping_scheme->GetThreadsPerBlock());
+
+ // TODO(b/110211620): Enable int32 index type for column reduction.
+ llvm::Type* index_ty =
+ is_column_reduction
+ ? b_.getInt64Ty()
+ : GetIndexTypeForKernel(unnested_hlo,
+ launch_dimensions.launch_bound(), &b_);
+
auto index_typed_constant = [&](uint64 c) -> llvm::Constant* {
return llvm::ConstantInt::get(index_ty, c);
};
@@ -2854,14 +2865,13 @@
// but we do it at the beginning in the hopes of reducing register pressure,
// since we touch threadIdx.x and blockIdx.x at the beginning of the kernel
// *anyway*.
- if (unnested_hlo->IsMultiOutputFusion()) {
- TF_CHECK_OK(KernelSupportLibrary(&b_).If(
+ if (!reduction_info && unnested_hlo->IsMultiOutputFusion()) {
+ KernelSupportLibrary{&b_}.If(
"emit_mof_tuple", IsBlock0Thread0(&b_), [&] {
llvm_ir::EmitTuple(GetIrArray(*unnested_hlo, *unnested_hlo),
ConstructIrArrayForOutputs(*unnested_hlo), &b_,
module_);
- return Status::OK();
- }));
+ });
}
// For each tiled parameter, cast its input IrArray to the corresponding
@@ -3243,15 +3253,17 @@
return std::make_tuple(num_reduced_major, num_kept, num_reduced_minor);
}
-std::tuple<KernelMappingScheme, bool> ComputeMappingSchemeAndReductionKind(
- const HloInstruction* first_reduce, llvm::IRBuilder<>* b) {
+} // namespace
+
+std::tuple<KernelMappingScheme, bool>
+IrEmitterUnnested::ComputeMappingSchemeAndReductionKind(
+ const HloInstruction* first_reduce) {
int64 depth = 1;
int64 height = 1;
int64 width = 1;
bool is_row_reduction = true;
int64 tile_size_x = 1;
int64 tile_size_y = 1;
- int64 block_size_y = 1;
int64 block_size_z = 1;
int64 num_threads_x = 1;
int64 num_threads_y = 1;
@@ -3274,14 +3286,17 @@
height = num_reduced_major;
width = num_kept;
is_row_reduction = false;
- tile_size_x = std::min(kWarpSize, num_kept);
- // The old Column reduction algorithm uses kTileHeight = 128. We choose
- // tile_size_y * block_size_y = 128 to match the value of kTileHeight. Using
- // a non-trivial block_size_y here is a way to avoid unrolling all the 128
- // iterations.
- tile_size_y = 32;
- block_size_y = 4;
+ // Column reduction without transpose doesn't require communication among
+ // threads processing elements in the same tile. The current implementation
+ // only support the use of on hardware thread block to process one block of
+ // tiles in the KernelMappingScheme. We try to maximize the values of
+ // num_threads_x and tile_size_x to allow a bigger hardware thread block.
+ int64 hw_threads_per_block_limit =
+ ThreadsPerBlockLimit(ir_emitter_context_->device_description());
+ tile_size_x = std::min(hw_threads_per_block_limit, num_kept);
num_threads_x = tile_size_x;
+ int64 kNumElementsPerPartialSum = 128;
+ tile_size_y = kNumElementsPerPartialSum;
} else {
// Row reduction reduces inputs with dimension [depth, height, width],
// where width is the most minor dimension, to dimension [height] .
@@ -3304,15 +3319,13 @@
<< " " << width;
DimensionVector dims_in_elem{depth, height, width};
- DimensionVector req_block_sizes{block_size_z, block_size_y, 1};
- llvm_ir::KernelMappingScheme mapping_scheme(dims_in_elem, tile_size_y,
- tile_size_x, req_block_sizes,
- num_threads_y, num_threads_x, b);
+ DimensionVector req_block_sizes{block_size_z, 1, 1};
+ llvm_ir::KernelMappingScheme mapping_scheme(
+ dims_in_elem, tile_size_y, tile_size_x, req_block_sizes, num_threads_y,
+ num_threads_x, &b_);
return std::make_tuple(mapping_scheme, is_row_reduction);
}
-} // namespace
-
Status IrEmitterUnnested::EmitReductionToVector(HloInstruction* unnested_hlo) {
VLOG(10) << "Emitting reduction to vector " << unnested_hlo->ToString();
@@ -3358,7 +3371,7 @@
bool is_row_reduction;
llvm_ir::KernelMappingScheme mapping_scheme;
std::tie(mapping_scheme, is_row_reduction) =
- ComputeMappingSchemeAndReductionKind(first_reduce, &b_);
+ ComputeMappingSchemeAndReductionKind(first_reduce);
ReductionCodegenInfo reduction_info(&mapping_scheme, is_row_reduction);
KernelCodeGenerator kernel_generator(
/*tile_element_generator=*/
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
index 85a0e53..1ebea7a 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
@@ -215,6 +215,11 @@
// Prerequisite: `IsReductionToVector(*unnested_hlo)`
Status EmitReductionToVector(HloInstruction* unnested_hlo);
+ // Computes the KernelMappingScheme for the reduce HLO and indicates whether
+ // the reduction is a row reduction.
+ std::tuple<llvm_ir::KernelMappingScheme, bool>
+ ComputeMappingSchemeAndReductionKind(const HloInstruction* first_reduce);
+
// Emits code for an in-place scatter, modifying `thunk`s launch dimensions in
// the process. `scatter` may be fused, scatter indices are taken from
// `scatter_indices_gen`, updates from`updates_gen`. The output buffer is
@@ -272,9 +277,8 @@
// For each reducer, emits the shuffle-down loop to accumulate the partial
// result to the global result.
void EmitFullWarpShuffleDownLoopForAllReduces(
- const absl::InlinedVector<HloComputation*, 1>& reducers,
- const absl::InlinedVector<llvm::AllocaInst*, 1>&
- partial_result_addresses);
+ absl::Span<HloComputation* const> reducers,
+ absl::Span<llvm::AllocaInst* const> partial_result_addresses);
// Generates the IrArray for each input of an hlo and returns a vector that
// constains such IrArrays.
diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc
index 24f07e6..bd53b90 100644
--- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc
+++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc
@@ -199,8 +199,7 @@
}
return absl::WrapUnique(target->createTargetMachine(
triple.str(), llvm_ir::AsStringRef(cpu_name), "+ptx60", target_options,
- Optional<Reloc::Model>(RelocModel), Optional<CodeModel::Model>(CMModel),
- codegen_opt_level));
+ getRelocModel(), getCodeModel(), codegen_opt_level));
}
// Adds the standard LLVM optimization passes, based on the speed optimization
@@ -394,8 +393,16 @@
int32 opt_level =
hlo_module_config.debug_options().xla_backend_optimization_level();
- CHECK_GE(opt_level, 2)
- << "The XLA GPU backend doesn't support unoptimized code generation";
+ if (opt_level < 2) {
+ LOG(ERROR) << std::string(80, '*');
+ LOG(ERROR) << "The XLA GPU backend doesn't support unoptimized code "
+ "generation but ";
+ LOG(ERROR) << "--xla_backend_optimization_level is set to " << opt_level
+ << "!";
+ LOG(ERROR) << "(Supported configuration is "
+ "--xla_backend_optimization_level >= 2.)";
+ LOG(ERROR) << std::string(80, '*');
+ }
AddOptimizationPasses(opt_level,
/*size_level=*/0, target_machine.get(), &module_passes,
diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
index e934cbd..60f2116 100644
--- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
+++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
@@ -108,27 +108,33 @@
namespace tracing = tensorflow::tracing;
-// Returns the directory containing nvvm libdevice files. config_cuda_data_dir
-// should be equal to config().debug_options().xla_gpu_cuda_data_dir() of the
-// HloModule being compiled.
-string GetLibdeviceDir(const string& config_cuda_data_dir) {
- std::vector<string> potential_libdevice_dirs;
- if (!config_cuda_data_dir.empty()) {
- potential_libdevice_dirs.push_back(config_cuda_data_dir);
- }
- potential_libdevice_dirs.push_back(tensorflow::LibdeviceRoot());
+// Returns a vector of potential locations of the CUDA root directory.
+std::vector<string> GetCudaRootCandidates(
+ const HloModuleConfig& hlo_module_config) {
+ std::vector<string> potential_cuda_roots = tensorflow::CandidateCudaRoots();
- // Tries all potential libdevice directories in the order they are inserted.
- // Returns the first directory that exists in the file system.
- for (const string& potential_libdevice_dir : potential_libdevice_dirs) {
- if (tensorflow::Env::Default()->IsDirectory(potential_libdevice_dir).ok()) {
- VLOG(2) << "Found libdevice dir " << potential_libdevice_dir;
- return potential_libdevice_dir;
+ // CUDA location explicitly specified by user via --xla_gpu_cuda_data_dir has
+ // highest priority.
+ string xla_gpu_cuda_data_dir =
+ hlo_module_config.debug_options().xla_gpu_cuda_data_dir();
+ if (!xla_gpu_cuda_data_dir.empty()) {
+ potential_cuda_roots.insert(potential_cuda_roots.begin(),
+ xla_gpu_cuda_data_dir);
+ }
+ return potential_cuda_roots;
+}
+
+// Returns the directory containing nvvm libdevice files.
+string GetLibdeviceDir(const HloModuleConfig& hlo_module_config) {
+ for (const string& cuda_root : GetCudaRootCandidates(hlo_module_config)) {
+ string libdevice_dir =
+ tensorflow::io::JoinPath(cuda_root, "nvvm", "libdevice");
+ VLOG(2) << "Looking for libdevice at " << libdevice_dir;
+ if (tensorflow::Env::Default()->IsDirectory(libdevice_dir).ok()) {
+ VLOG(2) << "Found libdevice dir " << libdevice_dir;
+ return libdevice_dir;
}
- VLOG(2) << "Unable to find potential libdevice dir "
- << potential_libdevice_dir;
}
-
LOG(WARNING) << "Unable to find libdevice dir. Using '.'";
// Last resort: maybe in the current folder.
return ".";
@@ -478,13 +484,19 @@
// Compiles the given PTX string using ptxas and returns the resulting machine
// code (i.e. a cubin) as a byte array.
-StatusOr<std::vector<uint8>> CompilePtx(const string& ptx, int cc_major,
- int cc_minor) {
+StatusOr<std::vector<uint8>> CompilePtx(
+ const string& ptx, int cc_major, int cc_minor,
+ const HloModuleConfig& hlo_module_config) {
tracing::ScopedActivity activity("Compile PTX", /*is_expensive=*/true);
- const string ptxas_path =
- tensorflow::io::JoinPath(tensorflow::CudaRoot(), "bin", "ptxas");
- VLOG(2) << "Checking ptxas at " << ptxas_path;
auto env = tensorflow::Env::Default();
+ string ptxas_path;
+ for (const string& cuda_root : GetCudaRootCandidates(hlo_module_config)) {
+ ptxas_path = tensorflow::io::JoinPath(cuda_root, "bin", "ptxas");
+ VLOG(2) << "Looking for ptxas at " << ptxas_path;
+ if (env->FileExists(ptxas_path).ok()) {
+ break;
+ }
+ }
TF_RETURN_IF_ERROR(env->FileExists(ptxas_path));
VLOG(2) << "Using ptxas at " << ptxas_path;
@@ -519,6 +531,9 @@
if (VLOG_IS_ON(2)) {
ptxas_args.push_back("-v");
}
+ if (hlo_module_config.debug_options().xla_gpu_disable_ptxas_optimizations()) {
+ ptxas_args.push_back("-O0");
+ }
ptxas_info_dumper.SetProgram(ptxas_path, ptxas_args);
ptxas_info_dumper.SetChannelAction(tensorflow::CHAN_STDERR,
tensorflow::ACTION_PIPE);
@@ -681,12 +696,8 @@
// Find the directory containing libdevice. To avoid searching for it every
// time, we have a one-element cache, keyed on the module's config's
// cuda_data_dir.
- const auto& config_cuda_data_dir =
- module->config().debug_options().xla_gpu_cuda_data_dir();
- if (cached_libdevice_dir_.empty() ||
- cached_cuda_data_dir_ != config_cuda_data_dir) {
- cached_cuda_data_dir_ = config_cuda_data_dir;
- cached_libdevice_dir_ = GetLibdeviceDir(config_cuda_data_dir);
+ if (cached_libdevice_dir_.empty()) {
+ cached_libdevice_dir_ = GetLibdeviceDir(module->config());
}
libdevice_dir = cached_libdevice_dir_;
}
@@ -740,7 +751,7 @@
}
const std::vector<uint8> cubin =
- CompilePtxOrGetCachedResult(ptx, cc_major, cc_minor);
+ CompilePtxOrGetCachedResult(ptx, cc_major, cc_minor, module->config());
auto thunk_schedule = absl::make_unique<ThunkSchedule>(
ir_emitter.ConsumeThunkSequence(), std::move(stream_assignment),
@@ -772,9 +783,9 @@
return std::unique_ptr<Executable>(gpu_executable);
}
-std::vector<uint8> NVPTXCompiler::CompilePtxOrGetCachedResult(const string& ptx,
- int cc_major,
- int cc_minor) {
+std::vector<uint8> NVPTXCompiler::CompilePtxOrGetCachedResult(
+ const string& ptx, int cc_major, int cc_minor,
+ const HloModuleConfig& hlo_module_config) {
XLA_SCOPED_LOGGING_TIMER("NVPTXCompiler::CompilePtxOrGetCachedResult");
tracing::ScopedActivity activity("PTX->CUBIN", /*is_expensive=*/true);
bool inserted;
@@ -803,7 +814,7 @@
CHECK(!cache_value->compilation_done);
if (!ptx.empty()) {
StatusOr<std::vector<uint8>> maybe_cubin =
- CompilePtx(*cache_ptx, cc_major, cc_minor);
+ CompilePtx(*cache_ptx, cc_major, cc_minor, hlo_module_config);
if (maybe_cubin.ok()) {
cache_value->cubin_data = std::move(maybe_cubin).ValueOrDie();
VLOG(2) << "Compiled PTX size:" << ptx.size()
diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h
index f79ae29..b2077f4 100644
--- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h
+++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h
@@ -97,8 +97,9 @@
// Tries to compile the given ptx string to cubin. Returns a vector with the
// compiled cubin. If compilation was unsuccessful, returns an empty vector.
- std::vector<uint8> CompilePtxOrGetCachedResult(const string& ptx,
- int cc_major, int cc_minor);
+ std::vector<uint8> CompilePtxOrGetCachedResult(
+ const string& ptx, int cc_major, int cc_minor,
+ const HloModuleConfig& hlo_module_config);
// The compilation_cache_ map is a cache from {ptx string, cc_major, cc_minor}
// -> cubin so we don't recompile the same ptx twice. This is important for
diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc
index 375f68a..bfed4f5 100644
--- a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc
+++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc
@@ -39,6 +39,25 @@
return out;
}
+int64 ThreadsPerBlockLimit(const se::DeviceDescription& device_desc) {
+ int64 threads_per_block = device_desc.threads_per_block_limit();
+ if (threads_per_block == 0) {
+ static std::atomic<int64> log_count{0};
+ if (log_count.fetch_add(1) < 8) {
+ LOG(WARNING) << "Attempting to calculate launch dimensions for GPU "
+ "without full information about its capabilities. "
+ "StreamExecutor's PopulateDeviceDescription should be "
+ "updated for this device.";
+ }
+ threads_per_block = device_desc.threads_per_warp();
+ if (threads_per_block == 0) {
+ // Fall back to *something* if we can't even get num threads per warp.
+ threads_per_block = 32;
+ }
+ }
+ return threads_per_block;
+}
+
// Calculates the launch dimensions used to invoke `hlo`.
LaunchDimensions CalculateLaunchDimensions(
const Shape& shape, const se::DeviceDescription& device_desc,
@@ -62,21 +81,7 @@
//
// <num threads per block> * <max blocks per core> = <max threads per core>
- int64 threads_per_block = device_desc.threads_per_block_limit();
- if (threads_per_block == 0) {
- static std::atomic<int64> log_count{0};
- if (log_count.fetch_add(1) < 8) {
- LOG(WARNING) << "Attempting to calculate launch dimensions for GPU "
- "without full information about its capabilities. "
- "StreamExecutor's PopulateDeviceDescription should be "
- "updated for this device.";
- }
- threads_per_block = device_desc.threads_per_warp();
- if (threads_per_block == 0) {
- // Fall back to *something* if we can't even get num threads per warp.
- threads_per_block = 32;
- }
- }
+ int64 threads_per_block = ThreadsPerBlockLimit(device_desc);
if (num_elements < threads_per_block) {
threads_per_block = num_elements;
diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.h b/tensorflow/compiler/xla/service/gpu/partition_assignment.h
index 0247112..eb41dcc 100644
--- a/tensorflow/compiler/xla/service/gpu/partition_assignment.h
+++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.h
@@ -57,6 +57,9 @@
std::ostream& operator<<(std::ostream& out,
const LaunchDimensions& launch_dims);
+// Returns the maximum number of threads per block allowed by the device.
+int64 ThreadsPerBlockLimit(const se::DeviceDescription& device_desc);
+
LaunchDimensions CalculateLaunchDimensions(
const Shape& shape, const se::DeviceDescription& device_desc,
int unroll_factor = 1);
diff --git a/tensorflow/compiler/xla/service/gpu/stream_executor_util.h b/tensorflow/compiler/xla/service/gpu/stream_executor_util.h
index 1fc46ba..92e4d6d 100644
--- a/tensorflow/compiler/xla/service/gpu/stream_executor_util.h
+++ b/tensorflow/compiler/xla/service/gpu/stream_executor_util.h
@@ -16,6 +16,7 @@
#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_EXECUTOR_UTIL_H_
#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_EXECUTOR_UTIL_H_
+#include "tensorflow/compiler/xla/layout.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index ff122b5..ca663b8 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -711,8 +711,6 @@
return eq(root_instruction(), other.root_instruction());
}
-uint64 HloComputation::Hash() const { return root_instruction()->Hash(); }
-
Status HloComputation::ReplaceWithNewInstruction(
HloInstruction* old_instruction,
std::unique_ptr<HloInstruction> new_instruction) {
@@ -797,7 +795,7 @@
template <typename HloInstructionPtr>
Status HloComputation::AcceptOrdered(
DfsHloVisitorBase<HloInstructionPtr>* visitor,
- const std::vector<HloInstruction*>& order) const {
+ absl::Span<HloInstruction* const> order) const {
VLOG(3) << "Accepting visitor with order.";
for (HloInstruction* root : CollectUnreachableRoots()) {
TF_RET_CHECK(std::find(order.begin(), order.end(), root) != order.end())
@@ -827,9 +825,9 @@
// Explicit instantiations.
template Status HloComputation::AcceptOrdered(
- DfsHloVisitor*, const std::vector<HloInstruction*>&) const;
+ DfsHloVisitor*, absl::Span<HloInstruction* const>) const;
template Status HloComputation::AcceptOrdered(
- ConstDfsHloVisitor*, const std::vector<HloInstruction*>&) const;
+ ConstDfsHloVisitor*, absl::Span<HloInstruction* const>) const;
Status HloComputation::Accept(
const std::function<Status(HloInstruction*)>& visitor_func) {
diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h
index c584e4c..5467d0a 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.h
+++ b/tensorflow/compiler/xla/service/hlo_computation.h
@@ -264,12 +264,6 @@
// Return whether `*this` and `other` are functionally equivalent.
bool operator==(const HloComputation& other) const;
- // Generates a hash value of an HLO computation. Hash considers
- // information on opcode, shape, operands, and typically a root instruction.
- // This function returns the same hash value for equivalent HLO computations,
- // with respect to HloInstruction::Identical() method.
- uint64 Hash() const;
-
// Replaces old instruction with newly created instruction. Removes old
// instruction from computation. Updates uses and root instruction.
Status ReplaceWithNewInstruction(
@@ -307,7 +301,7 @@
// be a topological sort of all instructions in the computation.
template <typename HloInstructionPtr>
Status AcceptOrdered(DfsHloVisitorBase<HloInstructionPtr>* visitor,
- const std::vector<HloInstruction*>& order) const;
+ absl::Span<HloInstruction* const> order) const;
// Same as Accept() above, but the visitor is given as a function.
Status Accept(const std::function<Status(HloInstruction*)>& visitor_func);
diff --git a/tensorflow/compiler/xla/service/hlo_computation_test.cc b/tensorflow/compiler/xla/service/hlo_computation_test.cc
index 8b50cfa..0361c87 100644
--- a/tensorflow/compiler/xla/service/hlo_computation_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation_test.cc
@@ -20,19 +20,19 @@
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
-#include "tensorflow/compiler/xla/service/hlo_matchers.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
+#include "tensorflow/compiler/xla/service/pattern_matcher.h"
+#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.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 op = xla::testing::opcode_matchers;
-
namespace xla {
namespace {
+namespace m = match;
using ::testing::ElementsAre;
using ::testing::UnorderedElementsAre;
@@ -261,7 +261,7 @@
auto computation = module->AddEntryComputation(builder.Build());
auto copy = computation->DeepCopyInstruction(constant).ValueOrDie();
- EXPECT_THAT(copy, op::Copy(constant));
+ EXPECT_THAT(copy, GmockMatch(m::Copy(m::Op().Is(constant))));
}
TEST_F(HloComputationTest, DeepCopyTuple) {
@@ -278,8 +278,9 @@
auto computation = module->AddEntryComputation(builder.Build());
auto tuple_copy = computation->DeepCopyInstruction(tuple).ValueOrDie();
- EXPECT_THAT(tuple_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)),
- op::Copy(op::GetTupleElement(tuple))));
+ EXPECT_THAT(tuple_copy, GmockMatch(m::Tuple(
+ m::Copy(m::GetTupleElement(m::Op().Is(tuple))),
+ m::Copy(m::GetTupleElement(m::Op().Is(tuple))))));
EXPECT_EQ(0, tuple_copy->operand(0)->operand(0)->tuple_index());
EXPECT_EQ(1, tuple_copy->operand(1)->operand(0)->tuple_index());
}
@@ -297,7 +298,7 @@
ShapeTree<bool> indices_to_copy(constant->shape(), /*init_value=*/true);
EXPECT_THAT(computation->DeepCopyInstruction(constant, &indices_to_copy)
.ValueOrDie(),
- op::Copy(constant));
+ GmockMatch(m::Copy(m::Op().Is(constant))));
}
{
@@ -330,10 +331,11 @@
computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added)
.ValueOrDie();
- EXPECT_THAT(deep_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)),
- op::Copy(op::GetTupleElement(tuple))));
- EXPECT_THAT(deep_copy, op::Tuple(copies_added.element({0}),
- copies_added.element({1})));
+ EXPECT_THAT(deep_copy, GmockMatch(m::Tuple(
+ m::Copy(m::GetTupleElement(m::Op().Is(tuple)))
+ .Is(copies_added.element({0})),
+ m::Copy(m::GetTupleElement(m::Op().Is(tuple)))
+ .Is(copies_added.element({1})))));
}
{
@@ -346,8 +348,9 @@
computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added)
.ValueOrDie();
- EXPECT_THAT(deep_copy, op::Tuple(op::GetTupleElement(tuple),
- op::GetTupleElement(tuple)));
+ EXPECT_THAT(deep_copy,
+ GmockMatch(m::Tuple(m::GetTupleElement(m::Op().Is(tuple)),
+ m::GetTupleElement(m::Op().Is(tuple)))));
EXPECT_TRUE(copies_added.element({}) == nullptr);
EXPECT_TRUE(copies_added.element({0}) == nullptr);
EXPECT_TRUE(copies_added.element({1}) == nullptr);
@@ -363,8 +366,9 @@
computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added)
.ValueOrDie();
- EXPECT_THAT(deep_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)),
- op::GetTupleElement(tuple)));
+ EXPECT_THAT(deep_copy, GmockMatch(m::Tuple(
+ m::Copy(m::GetTupleElement(m::Op().Is(tuple))),
+ m::GetTupleElement(m::Op().Is(tuple)))));
EXPECT_TRUE(copies_added.element({}) == nullptr);
EXPECT_TRUE(copies_added.element({0}) != nullptr);
EXPECT_TRUE(copies_added.element({1}) == nullptr);
@@ -381,7 +385,7 @@
auto copy = computation->DeepCopyInstruction(token).ValueOrDie();
// No copy should be added.
- EXPECT_THAT(copy, op::AfterAll());
+ EXPECT_THAT(copy, GmockMatch(m::AfterAll()));
}
TEST_F(HloComputationTest, DeepCopyTokenTuple) {
@@ -399,8 +403,9 @@
// Only the array (second tuple element) should be copied. The token is passed
// through transparently.
- EXPECT_THAT(copy, op::Tuple(op::GetTupleElement(tuple),
- op::Copy(op::GetTupleElement(tuple))));
+ EXPECT_THAT(copy, GmockMatch(m::Tuple(
+ m::GetTupleElement(m::Op().Is(tuple)),
+ m::Copy(m::GetTupleElement(m::Op().Is(tuple))))));
}
TEST_F(HloComputationTest, CycleDetection) {
@@ -443,13 +448,15 @@
auto module = CreateNewVerifiedModule();
auto computation = module->AddEntryComputation(builder.Build());
EXPECT_EQ(4, computation->instruction_count());
- EXPECT_THAT(computation->root_instruction(), op::Negate(constant));
+ EXPECT_THAT(computation->root_instruction(),
+ GmockMatch(m::Negate(m::Op().Is(constant))));
EXPECT_EQ(negate, computation->root_instruction());
ASSERT_IS_OK(computation->RemoveInstructionAndUnusedOperands(dead_add));
EXPECT_EQ(2, computation->instruction_count());
- EXPECT_THAT(computation->root_instruction(), op::Negate(constant));
+ EXPECT_THAT(computation->root_instruction(),
+ GmockMatch(m::Negate(m::Op().Is(constant))));
EXPECT_EQ(negate, computation->root_instruction());
}
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc
index 3a7652a..934c082 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator.cc
+++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc
@@ -33,6 +33,7 @@
#include "tensorflow/compiler/xla/literal_util.h"
#include "tensorflow/compiler/xla/map_util.h"
#include "tensorflow/compiler/xla/primitive_util.h"
+#include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h"
#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
@@ -629,8 +630,11 @@
evaluated_[compare],
Compare<int64>(compare->shape(), opcode, lhs_literal, rhs_literal));
} break;
- case F16:
- return Unimplemented("unhandled primitive type: F16.");
+ case F16: {
+ TF_ASSIGN_OR_RETURN(
+ evaluated_[compare],
+ Compare<half>(compare->shape(), opcode, lhs_literal, rhs_literal));
+ } break;
case BF16: {
TF_ASSIGN_OR_RETURN(evaluated_[compare],
Compare<bfloat16>(compare->shape(), opcode,
@@ -1449,4 +1453,46 @@
template StatusOr<Literal> HloEvaluator::Evaluate<const Literal*>(
HloInstruction* instruction, absl::Span<const Literal* const> arg_literals);
+namespace {
+template <typename T>
+std::unique_ptr<Array2D<T>> MatmulArray2DImpl(
+ const Array2D<T>& lhs, const Array2D<T>& rhs,
+ const std::function<void(
+ const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, int64 n,
+ int64 k, int32 transpose_lhs, int32 transpose_rhs)>& impl_fn) {
+ CHECK_EQ(lhs.width(), rhs.height());
+ int m = lhs.height();
+ int n = rhs.width();
+ int k = lhs.width();
+ auto result = absl::make_unique<Array2D<T>>(m, n);
+ // Because Eigen is a header-oriented library, make sure that the Eigen code
+ // is the same as the code used by the CPU backend (otherwise the linker will
+ // randomly pick *some* definition).
+ impl_fn(
+ /*run_options_ptr=*/nullptr, result->data(), rhs.data(), lhs.data(), n, m,
+ k,
+ /*transpose_lhs=*/0,
+ /*transpose_rhs=*/0);
+ return result;
+}
+} // namespace
+
+std::unique_ptr<Array2D<Eigen::half>> HloEvaluator::MatmulArray2D(
+ const Array2D<Eigen::half>& lhs, const Array2D<Eigen::half>& rhs) {
+ return MatmulArray2DImpl<Eigen::half>(
+ lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF16);
+}
+
+std::unique_ptr<Array2D<float>> HloEvaluator::MatmulArray2D(
+ const Array2D<float>& lhs, const Array2D<float>& rhs) {
+ return MatmulArray2DImpl<float>(
+ lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF32);
+}
+
+std::unique_ptr<Array2D<double>> HloEvaluator::MatmulArray2D(
+ const Array2D<double>& lhs, const Array2D<double>& rhs) {
+ return MatmulArray2DImpl<double>(
+ lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF64);
+}
+
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h
index 45ed813..d363a51 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator.h
+++ b/tensorflow/compiler/xla/service/hlo_evaluator.h
@@ -21,6 +21,7 @@
#include "absl/container/node_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/types/span.h"
+#include "tensorflow/compiler/xla/array2d.h"
#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
@@ -119,6 +120,17 @@
const PrecisionConfig& precision_config,
const Literal& lhs, const Literal& rhs);
+ // Enable the fast path for certain operations like dot or convolution.
+ void set_use_fast_path(bool value) { use_fast_path_ = value; }
+
+ // Returns the result of a matrix multiply `lhs x rhs`.
+ static std::unique_ptr<Array2D<Eigen::half>> MatmulArray2D(
+ const Array2D<Eigen::half>& lhs, const Array2D<Eigen::half>& rhs);
+ static std::unique_ptr<Array2D<float>> MatmulArray2D(
+ const Array2D<float>& lhs, const Array2D<float>& rhs);
+ static std::unique_ptr<Array2D<double>> MatmulArray2D(
+ const Array2D<double>& lhs, const Array2D<double>& rhs);
+
protected:
// Make HloEvaluatorTypedVisitor a friend because it is logically part of this
// class.
@@ -217,6 +229,9 @@
// we cannot use flat_hash_map any more.
absl::node_hash_map<const HloInstruction*, Literal> evaluated_;
+ // Use fast path that uses eigen in the evaluator.
+ bool use_fast_path_ = false;
+
private:
template <typename ReturnT, typename NativeT>
static StatusOr<Literal> ElementWiseUnaryOpImpl(
@@ -250,6 +265,8 @@
TF_DISALLOW_COPY_AND_ASSIGN(HloEvaluator);
};
+std::unique_ptr<Array2D<float>> MatmulArray2D(const Array2D<float>& lhs,
+ const Array2D<float>& rhs);
} // namespace xla
#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h
index b87fc3e..03d4299 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h
@@ -23,6 +23,7 @@
#include "absl/container/inlined_vector.h"
#include "absl/memory/memory.h"
#include "absl/types/optional.h"
+#include "tensorflow/compiler/xla/array2d.h"
#include "tensorflow/compiler/xla/literal_util.h"
#include "tensorflow/compiler/xla/service/hlo_casting_utils.h"
#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
@@ -105,6 +106,12 @@
template <typename ReturnT, typename ElementwiseT = ReturnT>
class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault {
private:
+ Status UnsupportedTypeError(HloInstruction* instruction) {
+ return InvalidArgument(
+ "Unsupported type for %s: %s", HloOpcodeString(instruction->opcode()),
+ PrimitiveType_Name(instruction->shape().element_type()));
+ }
+
// Get the value in the given literal static_cast as a double.
template <
typename NativeT,
@@ -224,7 +231,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleRound(HloInstruction* round) {
- return InvalidArgument("Unsupported type for Round");
+ return UnsupportedTypeError(round);
}
Status HandleRound(HloInstruction* round) override {
@@ -246,7 +253,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleCeil(HloInstruction* ceil) {
- return InvalidArgument("Unsupported type for Ceil");
+ return UnsupportedTypeError(ceil);
}
Status HandleCeil(HloInstruction* ceil) override {
@@ -297,8 +304,8 @@
template <
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
- Status HandleExpm1(HloInstruction* floor) {
- return InvalidArgument("Unsupported type for Expm1");
+ Status HandleExpm1(HloInstruction* expm1) {
+ return UnsupportedTypeError(expm1);
}
Status HandleExpm1(HloInstruction* floor) override {
@@ -321,7 +328,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleFloor(HloInstruction* floor) {
- return InvalidArgument("Unsupported type for Floor");
+ return UnsupportedTypeError(floor);
}
Status HandleFloor(HloInstruction* floor) override {
@@ -351,12 +358,12 @@
template <
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
- Status HandleLog1p(HloInstruction* floor) {
- return InvalidArgument("Unsupported type for Log1p");
+ Status HandleLog1p(HloInstruction* log1p) {
+ return UnsupportedTypeError(log1p);
}
- Status HandleLog1p(HloInstruction* floor) override {
- return HandleLog1p<ReturnT>(floor);
+ Status HandleLog1p(HloInstruction* log1p) override {
+ return HandleLog1p<ReturnT>(log1p);
}
template <typename NativeT,
@@ -396,7 +403,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleNot(HloInstruction* not_) {
- return InvalidArgument("Unsupported type for Not");
+ return UnsupportedTypeError(not_);
}
Status HandleNot(HloInstruction* not_) override {
@@ -476,7 +483,7 @@
template <typename NativeT, typename std::enable_if<!std::is_floating_point<
NativeT>::value>::type* = nullptr>
Status HandleAtan2(HloInstruction* atan2) {
- return InvalidArgument("Unsupported type for Atan2");
+ return UnsupportedTypeError(atan2);
}
Status HandleAtan2(HloInstruction* atan2) override {
@@ -624,7 +631,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleMaximum(HloInstruction* maximum) {
- return InvalidArgument("Unsupported type for Maximum");
+ return UnsupportedTypeError(maximum);
}
Status HandleMaximum(HloInstruction* maximum) override {
@@ -659,7 +666,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleMinimum(HloInstruction* minimum) {
- return InvalidArgument("Unsupported type for Minimum");
+ return UnsupportedTypeError(minimum);
}
Status HandleMinimum(HloInstruction* minimum) override {
@@ -724,7 +731,7 @@
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleRemainder(HloInstruction* remainder) {
- return InvalidArgument("Unsupported type for Remainder");
+ return UnsupportedTypeError(remainder);
}
Status HandleRemainder(HloInstruction* remainder) override {
@@ -746,14 +753,14 @@
template <typename NativeT, typename std::enable_if<std::is_floating_point<
NativeT>::value>::type* = nullptr>
Status HandleAnd(HloInstruction* and_) {
- return InvalidArgument("Unsupported type for And");
+ return UnsupportedTypeError(and_);
}
template <
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleAnd(HloInstruction* and_) {
- return InvalidArgument("Unsupported type for And");
+ return UnsupportedTypeError(and_);
}
Status HandleAnd(HloInstruction* and_) override {
@@ -775,7 +782,7 @@
template <typename NativeT, typename std::enable_if<std::is_floating_point<
NativeT>::value>::type* = nullptr>
Status HandleOr(HloInstruction* or_) {
- return InvalidArgument("Unsupported type for Or");
+ return UnsupportedTypeError(or_);
}
template <
@@ -804,14 +811,14 @@
template <typename NativeT, typename std::enable_if<std::is_floating_point<
NativeT>::value>::type* = nullptr>
Status HandleXor(HloInstruction* xor_) {
- return InvalidArgument("Unsupported type for Xor");
+ return UnsupportedTypeError(xor_);
}
template <
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleXor(HloInstruction* xor_) {
- return InvalidArgument("Unsupported type for Xor");
+ return UnsupportedTypeError(xor_);
}
Status HandleXor(HloInstruction* xor_) override {
@@ -836,8 +843,8 @@
typename std::enable_if<!std::is_integral<NativeT>::value ||
std::is_same<NativeT, bool>::value>::type* =
nullptr>
- Status HandleShiftLeft(HloInstruction*) {
- return InvalidArgument("Unsupported type for ShiftLeft");
+ Status HandleShiftLeft(HloInstruction* shift) {
+ return UnsupportedTypeError(shift);
}
Status HandleShiftLeft(HloInstruction* shl) override {
@@ -866,8 +873,8 @@
typename std::enable_if<!std::is_integral<NativeT>::value ||
std::is_same<NativeT, bool>::value>::type* =
nullptr>
- Status HandleShiftRightArithmetic(HloInstruction*) {
- return InvalidArgument("Unsupported type for ShiftRightArithmetic");
+ Status HandleShiftRightArithmetic(HloInstruction* shift) {
+ return UnsupportedTypeError(shift);
}
Status HandleShiftRightArithmetic(HloInstruction* shra) override {
@@ -897,8 +904,8 @@
typename std::enable_if<!std::is_integral<NativeT>::value ||
std::is_same<NativeT, bool>::value>::type* =
nullptr>
- Status HandleShiftRightLogical(HloInstruction*) {
- return InvalidArgument("Unsupported type for ShiftRightLogical");
+ Status HandleShiftRightLogical(HloInstruction* shift) {
+ return UnsupportedTypeError(shift);
}
Status HandleShiftRightLogical(HloInstruction* shrl) override {
@@ -923,8 +930,8 @@
template <
typename NativeT,
typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
- Status HandleClamp(HloInstruction*) {
- return InvalidArgument("Unsupported type for Clamp");
+ Status HandleClamp(HloInstruction* clamp) {
+ return UnsupportedTypeError(clamp);
}
Status HandleClamp(HloInstruction* clamp) override {
@@ -1148,6 +1155,78 @@
}
Status HandleDot(HloInstruction* dot) override {
+ if (parent_->use_fast_path_) {
+ return HandleDot<ReturnT>(dot);
+ }
+ return HandleDotSlowPath(dot);
+ }
+
+ template <typename NativeT, typename std::enable_if<std::is_same<
+ NativeT, float>::value>::type* = nullptr>
+ Status HandleDot(HloInstruction* dot) {
+ const HloInstruction* lhs = dot->operand(0);
+ const HloInstruction* rhs = dot->operand(1);
+ CHECK(ShapeUtil::IsArray(dot->shape()));
+ CHECK(ShapeUtil::IsArray(lhs->shape()));
+ CHECK(ShapeUtil::IsArray(rhs->shape()));
+
+ const auto& dnums = dot->dot_dimension_numbers();
+
+ const int64 lhs_rank = ShapeUtil::Rank(lhs->shape());
+ const int64 rhs_rank = ShapeUtil::Rank(rhs->shape());
+
+ CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape()));
+ CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape()));
+
+ // There must be 1 and only 1 Contracting dimension for lhs and rhs.
+ CHECK_EQ(dnums.lhs_contracting_dimensions_size(), 1);
+ CHECK_EQ(dnums.rhs_contracting_dimensions_size(), 1);
+ const int64 lhs_contracting_dimension = dnums.lhs_contracting_dimensions(0);
+ const int64 rhs_contracting_dimension = dnums.rhs_contracting_dimensions(0);
+ // Contracted dimension sizes must be the same.
+ CHECK_EQ(lhs->shape().dimensions(lhs_contracting_dimension),
+ rhs->shape().dimensions(rhs_contracting_dimension))
+ << "lhs contracted dimension: "
+ << lhs->shape().dimensions(lhs_contracting_dimension)
+ << " rhs contracted dimension: "
+ << rhs->shape().dimensions(rhs_contracting_dimension);
+
+ // The fast path is for a simple rank 2 dot with default layout operands.
+ if (lhs_rank == 2 && rhs_rank == 2 && lhs_contracting_dimension == 1 &&
+ rhs_contracting_dimension == 0 &&
+ LayoutUtil::Equal(lhs->shape().layout(),
+ LayoutUtil::GetDefaultLayoutForR2()) &&
+ LayoutUtil::Equal(rhs->shape().layout(),
+ LayoutUtil::GetDefaultLayoutForR2()) &&
+ LayoutUtil::Equal(dot->shape().layout(),
+ LayoutUtil::GetDefaultLayoutForR2())) {
+ const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
+ const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
+ const int64 contracted_dimension_size =
+ lhs->shape().dimensions(lhs_contracting_dimension);
+ Array2D<NativeT> lhs_array(lhs->shape().dimensions(0),
+ contracted_dimension_size);
+ lhs_array.SetValues(lhs_literal.data<NativeT>());
+ Array2D<NativeT> rhs_array(contracted_dimension_size,
+ rhs->shape().dimensions(1));
+ rhs_array.SetValues(rhs_literal.data<NativeT>());
+ std::unique_ptr<Array2D<NativeT>> result_array =
+ HloEvaluator::MatmulArray2D(lhs_array, rhs_array);
+ Literal result(dot->shape());
+ result.PopulateR2FromArray2D(*result_array);
+ parent_->evaluated_[dot] = std::move(result);
+ return Status::OK();
+ }
+ return HandleDotSlowPath(dot);
+ }
+
+ template <typename NativeT, typename std::enable_if<!std::is_same<
+ NativeT, float>::value>::type* = nullptr>
+ Status HandleDot(HloInstruction* dot) {
+ return HandleDotSlowPath(dot);
+ }
+
+ Status HandleDotSlowPath(HloInstruction* dot) {
auto lhs = dot->operand(0);
auto rhs = dot->operand(1);
CHECK(ShapeUtil::IsArray(dot->shape()));
@@ -1578,7 +1657,7 @@
std::is_same<NativeT, bool>::value>::type* =
nullptr>
Status HandleSort(HloInstruction* sort) {
- return InvalidArgument("Unsupported type for Sort");
+ return UnsupportedTypeError(sort);
}
Status HandleSort(HloInstruction* sort) override {
@@ -2357,7 +2436,7 @@
std::is_same<NativeT, int64>::value ||
std::is_same<NativeT, uint64>::value)>::type* = nullptr>
Status HandleClz(HloInstruction* clz) {
- return InvalidArgument("Unsupported type for Clz");
+ return UnsupportedTypeError(clz);
}
template <typename NativeT,
@@ -2403,7 +2482,7 @@
typename std::enable_if<std::is_integral<NativeT>::value ||
is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleSin(HloInstruction* sin) {
- return InvalidArgument("Unsupported type for Sin");
+ return UnsupportedTypeError(sin);
}
Status HandleSin(HloInstruction* sin) override {
@@ -2425,7 +2504,7 @@
typename std::enable_if<std::is_integral<NativeT>::value ||
is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleCos(HloInstruction* cos) {
- return InvalidArgument("Unsupported type for Cos");
+ return UnsupportedTypeError(cos);
}
Status HandleCos(HloInstruction* cos) override {
@@ -2534,7 +2613,7 @@
typename std::enable_if<std::is_integral<NativeT>::value ||
is_complex_t<NativeT>::value>::type* = nullptr>
Status HandleReducePrecision(HloInstruction* reduce_precision) {
- return InvalidArgument("Unsupported type for reduce precision");
+ return UnsupportedTypeError(reduce_precision);
}
Status HandleReducePrecision(HloInstruction* reduce_precision) override {
@@ -2543,15 +2622,27 @@
template <typename NativeT,
typename std::enable_if<
+ std::is_same<NativeT, bfloat16>::value ||
+ std::is_same<NativeT, Eigen::half>::value ||
std::is_integral<NativeT>::value ||
std::is_floating_point<NativeT>::value>::type* = nullptr>
Status HandleIota(HloInstruction* instruction) {
auto* iota = Cast<HloIotaInstruction>(instruction);
+ const int64 iota_size = iota->shape().dimensions(iota->iota_dimension());
// Avoid using std::vector since std::vector<bool> does not convert to
// absl::Span<bool>.
- absl::InlinedVector<NativeT, 1> data(
- iota->shape().dimensions(iota->iota_dimension()));
- std::iota(data.begin(), data.end(), 0);
+ absl::InlinedVector<NativeT, 1> data(iota_size);
+ // We don't use std::iota for two reasons:
+ //
+ // (1) std:iota does not support bfloat16 and float16.
+ //
+ // (2) std::iota saturates for floating point types when the value is not
+ // representable, but the definition of HLO iota is the value as a
+ // 64-bit integer cast to the native type.
+ for (int64 i = 0; i < iota_size; ++i) {
+ // static_cast is required for Eigen::half (F16).
+ data[i] = static_cast<NativeT>(i);
+ }
auto result = LiteralUtil::CreateR1<NativeT>(data);
if (ShapeUtil::Rank(iota->shape()) > 1) {
@@ -2567,10 +2658,12 @@
}
template <typename NativeT,
typename std::enable_if<
- !(std::is_integral<NativeT>::value ||
+ !(std::is_same<NativeT, bfloat16>::value ||
+ std::is_same<NativeT, Eigen::half>::value ||
+ std::is_integral<NativeT>::value ||
std::is_floating_point<NativeT>::value)>::type* = nullptr>
Status HandleIota(HloInstruction* iota) {
- return InvalidArgument("Unsupported type for iota");
+ return UnsupportedTypeError(iota);
}
Status HandleIota(HloInstruction* iota) override {
return HandleIota<ReturnT>(iota);
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
index 302eca6..5db21e4 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
@@ -1474,14 +1474,15 @@
GraphRendererInterface::GraphKind graph_kind,
const DebugOptions& debug_options) {
string path = debug_options.xla_hlo_graph_path();
- if (!path.empty()) {
+ if (!path.empty() && !debug_options.xla_hlo_dump_as_html()) {
return SaveGraph(graph, graph_kind, path);
} else {
auto graph_renderer =
GraphRendererRegistry::Default()->GetDefaultRenderer();
CHECK(graph_renderer != nullptr)
<< "No registered renderer for the HLO graph. "
- "Use --xla_hlo_graph_path=PATH to export to local file system";
+ "Use --xla_hlo_graph_path=PATH --xla_hlo_dump_as_html=false to "
+ "export to local file system";
return graph_renderer->RenderGraph(graph, graph_kind, debug_options);
}
}
@@ -1589,5 +1590,143 @@
return graph_url;
}
+string WrapDotInHTML(const string& dot) {
+ static const char html_prefix[] = R"html(
+<!DOCTYPE html>
+<html>
+<head>
+ <meta charset="utf-8">
+ <style type="text/css">
+ html, body { height: 100%; }
+ body { margin: 0; }
+ </style>
+</head>
+<body>
+ <!-- Integrity hash is generated by https://www.srihash.org/ -->
+ <script src="https://cdn.jsdelivr.net/npm/viz.js@2.1.1/viz.js"
+ integrity="sha384-aD1MJYb0WKIUT+CtwJp5LTuV3U4pLAS6B/nUxL7ECimC2pN9N8vjlMr/yQCAkzxE"
+ crossorigin="anonymous"></script>
+ <script src="https://cdn.jsdelivr.net/npm/viz.js@2.1.1/full.render.js"
+ integrity="sha384-bAixY275aIpCj6Te19y0MILZ4V+VEC8CVFujFEH+Lf7W+4XYYeYLwW5IBI6yQmMT"
+ crossorigin="anonymous"></script>
+ <script src="https://cdn.jsdelivr.net/npm/svg-pan-zoom@3.6.0/dist/svg-pan-zoom.min.js"
+ integrity="sha384-3008WpYB2pOBvE7lwkrKf+qTmbTPGGPYxA9C1YVhvbPukns4ZFj7E98QPLkNW9dS"
+ crossorigin="anonymous"></script>
+ <div id="container" style="height:95%; border:1px solid black; "></div>
+ <script>
+ var data = `
+)html";
+
+ static const char html_suffix[] = R"html(
+`;
+ var cssregex = new RegExp('stylesheet=<([^]*)\n>\n', 'gm');
+ var results = cssregex.exec(data)
+ // graphviz has problem dealing with large stylesheets.
+ // https://github.com/tensorflow/tensorflow/issues/17220#issuecomment-369228492
+ // In order to avoid the problem, remove the stylesheet from the dot and
+ // insert it directly info the rendered SVG.
+ var dot_data = data;
+ var css_data = ''
+ if (results !== null) {
+ css_data = results[1].replace(/\s*data:.*\s*,/,''); // Strip content-type field.
+ dot_data = data.replace(cssregex, ''); // Remove the stylesheet
+ }
+
+ var render_start = performance.now()
+ function add_controls(svg) {
+ var htmlblob = new Blob([document.documentElement.innerHTML],
+ {type: 'text/html'});
+ var savehtml = document.createElement('a');
+ savehtml.setAttribute('href', URL.createObjectURL(htmlblob));
+ savehtml.setAttribute('download', 'graph.html');
+ savehtml.innerHTML = " [Save HTML+SVG] ";
+ document.body.append(savehtml);
+ var svgblob = new Blob([svg.outerHTML], {type: 'image/svg'});
+ var savesvg = document.createElement('a');
+ savesvg.setAttribute('href', URL.createObjectURL(svgblob));
+ savesvg.setAttribute('download', 'graph.svg');
+ savesvg.innerHTML = " [Save SVG] ";
+ document.body.append(savesvg);
+ var dotblob = new Blob([data], {type: 'text/dot'});
+ var savedot = document.createElement('a');
+ savedot.setAttribute('href', URL.createObjectURL(dotblob));
+ savedot.setAttribute('download', 'graph.dot');
+ savedot.innerHTML = " [Save DOT] ";
+ document.body.append(savedot);
+ // Will get called after embed element was loaded
+ var panzoom = svgPanZoom(svg, {
+ zoomEnabled: true,
+ controlIconsEnabled: true,
+ });
+ document.getElementsByTagName("BODY")[0].onresize = function() {
+ panzoom.resize();
+ panzoom.fit();
+ panzoom.center();
+ };
+ var render_end = performance.now();
+ var render_note = document.createElement('div')
+ render_note.innerHTML = 'Rendering took '
+ + (render_end - render_start).toFixed(2) + "ms."
+ document.body.append(render_note);
+ }
+ var svg = document.getElementById('graph')
+ if (svg == null) {
+ // Need to render SVG first.
+ var viz = new Viz();
+ viz.renderSVGElement(dot_data)
+ .then(function(svg){
+ var container = document.getElementById('container')
+ var style = document.createElementNS('http://www.w3.org/2000/svg', 'style');
+ var node = document.createTextNode(css_data);
+ style.appendChild(node);
+ svg.setAttribute('width', '100%');
+ svg.setAttribute('height', 'auto');
+ svg.setAttribute('id', 'graph');
+ svg.appendChild(style);
+ container.appendChild(svg);
+ add_controls(svg);
+ })
+ } else {
+ // HTML already has rendered SVG embedded, so we just need to add
+ // controls.
+ add_controls(svg);
+ }
+ </script>
+</body>
+</html>
+)html";
+
+ return html_prefix + dot + html_suffix;
+}
+
+string RenderDotAsHTMLFile(const string& dot,
+ const DebugOptions& debug_options) {
+ string html = WrapDotInHTML(dot);
+
+ auto env = tensorflow::Env::Default();
+ std::vector<string> dirs;
+ string output_dir = debug_options.xla_hlo_graph_path();
+ if (output_dir.empty()) {
+ env->GetLocalTempDirectories(&dirs);
+ } else {
+ dirs.push_back(output_dir);
+ }
+ // Try each directory, as they might be full, have inappropriate
+ // permissions or have different problems at times.
+ string output;
+ for (const string& dir : dirs) {
+ string filename = tensorflow::io::JoinPath(dir, "graph-");
+ if (env->CreateUniqueFileName(&filename, ".html")) {
+ output = filename;
+ break;
+ }
+ }
+ if (output.empty()) {
+ LOG(FATAL) << "Failed to create unique output file name.";
+ }
+ TF_CHECK_OK(tensorflow::WriteStringToFile(env, output, html));
+ return "file://" + output;
+}
+
} // namespace hlo_graph_dumper
} // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.h b/tensorflow/compiler/xla/service/hlo_graph_dumper.h
index de1eefa..8e51454 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper.h
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.h
@@ -81,6 +81,12 @@
void DumpText(const HloModule& module, const string& label,
const string& directory_path, bool do_prefix = true);
+// Renders DOT graph as inline SVG and saves it in an HTML file in a temprary
+// directory or directory specified via --xla_hlo_graph_path. Returns the file
+// URI pointing to the file.
+string RenderDotAsHTMLFile(const string& dot,
+ const DebugOptions& debug_options);
+
// Graph renderers may be added using a registration mechanism, e.g.:
// XLA_REGISTER_GRAPH_RENDERER(AGraphRendererClass, 100)
// The renderer with the highest numeric priority value is used.
diff --git a/tensorflow/compiler/xla/service/hlo_graph_html_renderer.cc b/tensorflow/compiler/xla/service/hlo_graph_html_renderer.cc
new file mode 100644
index 0000000..84c4cf1
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_graph_html_renderer.cc
@@ -0,0 +1,43 @@
+/* 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 of an DOT graph renderer that uses Javascript to render DOT to
+// SVG in a browser.
+
+#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+
+namespace xla {
+namespace hlo_graph_dumper {
+namespace {
+
+class GraphHtmlRenderer : public GraphRendererInterface {
+ public:
+ string RenderGraph(const string& graph, GraphKind graph_kind,
+ const DebugOptions& debug_options) override {
+ switch (graph_kind) {
+ case DOT_GRAPH:
+ return RenderDotAsHTMLFile(graph, debug_options);
+ default:
+ LOG(FATAL) << "Only DOT graphs can be rendered";
+ }
+ }
+};
+
+XLA_REGISTER_GRAPH_RENDERER(GraphHtmlRenderer);
+
+} // namespace
+} // namespace hlo_graph_dumper
+} // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index 21b1dbc..8b2ace1 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -569,6 +569,11 @@
instruction->SetAndSanitizeName(proto.name());
instruction->metadata_ = proto.metadata();
instruction->backend_config_ = proto.backend_config();
+
+ TF_RET_CHECK(proto.id() >= 0)
+ << "Instruction with negative id: " << proto.id();
+ TF_RET_CHECK(proto.id() <= INT_MAX)
+ << "Instruction with id > INT_MAX: " << proto.id();
instruction->unique_id_ = proto.id();
if (proto.has_sharding()) {
@@ -914,12 +919,8 @@
HloInstruction* operand,
HloInstruction* update,
HloInstruction* start_indices) {
- auto instruction = absl::WrapUnique(
- new HloInstruction(HloOpcode::kDynamicUpdateSlice, shape));
- instruction->AppendOperand(operand);
- instruction->AppendOperand(update);
- instruction->AppendOperand(start_indices);
- return instruction;
+ return absl::make_unique<HloDynamicUpdateSliceInstruction>(
+ shape, operand, update, start_indices);
}
/* static */ std::unique_ptr<HloInstruction> HloInstruction::CreateConcatenate(
@@ -1760,7 +1761,12 @@
return false;
}
-uint64 HloInstruction::Hash() const {
+static uint64 HashOperand(const HloInstruction* hlo) {
+ return ShapeUtil::Hash(hlo->shape());
+}
+
+uint64 HloInstruction::Hash(
+ const std::function<uint64(const HloInstruction*)>& hash_operand) const {
using tensorflow::Hash64Combine;
uint64 hash_value = Hash64Combine(0, static_cast<uint64>(opcode()));
@@ -1769,7 +1775,7 @@
if (!IsCrossModuleAllReduce()) {
if (!operands().empty()) {
for (size_t i = 0; i < operands().size(); ++i) {
- hash_value = Hash64Combine(hash_value, operand(i)->Hash());
+ hash_value = Hash64Combine(hash_value, hash_operand(operand(i)));
}
}
}
@@ -1778,6 +1784,11 @@
return hash_value;
}
+uint64 HloInstruction::Hash() const {
+ // Use HashOperand as an argument to prevent non-termination.
+ return Hash(HashOperand);
+}
+
uint64 HloInstruction::InnerHash() const { return 13; }
void HloInstruction::RemoveUser(HloInstruction* user) {
@@ -2059,6 +2070,10 @@
return opcode() == HloOpcode::kCrossReplicaSum && all_reduce_id();
}
+bool HloInstruction::IsCrossReplicaAllReduce() const {
+ return opcode() == HloOpcode::kCrossReplicaSum && !all_reduce_id();
+}
+
string HloInstruction::ToStringWithCanonicalNameMap(
const HloPrintOptions& options,
CanonicalNameMap* canonical_name_map) const {
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index a547162..dd77f10 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -909,6 +909,14 @@
// information on opcode, shape, operands, and typically a root instruction.
// This function returns the same hash value for equivalent HLO instructions,
// with respect to HloInstruction::Identical() method.
+ //
+ // Uses hash_operand function to compute hash values of its operands.
+ // At the very top level, hash_operand should be non-recursive to prevent
+ // non-termination.
+ uint64 Hash(
+ const std::function<uint64(const HloInstruction*)>& hash_operand) const;
+
+ // Calls the above method with non-recursive hash_operand function.
uint64 Hash() const;
// Returns whether the instruction has a constant operand.
@@ -1174,9 +1182,12 @@
// Returns true if this instruction is elementwise on all its operands.
bool IsElementwise() const;
- // Returns true if this is an cross module all-reduce instrucion.
+ // Returns true if this is a cross module all-reduce instruction.
bool IsCrossModuleAllReduce() const;
+ // Returns true if this is a cross-replica all-reduce instruction.
+ bool IsCrossReplicaAllReduce() const;
+
// Returns true if this elementwise instruction implicitly broadcasts operand
// `operand_idx`.
//
diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc
index 1ea02cf..f55de6a 100644
--- a/tensorflow/compiler/xla/service/hlo_instructions.cc
+++ b/tensorflow/compiler/xla/service/hlo_instructions.cc
@@ -1372,8 +1372,14 @@
other.fused_instructions_computation());
}
+static uint64 HashOperandRecursive(const HloInstruction* hlo) {
+ return hlo->Hash(HashOperandRecursive);
+}
+
uint64 HloFusionInstruction::InnerHash() const {
- return fused_instructions_computation()->Hash();
+ // Use HashOperandRecursive to recursively compute hash on inner operands.
+ return fused_instructions_computation()->root_instruction()->Hash(
+ HashOperandRecursive);
}
std::unique_ptr<HloInstruction> HloFusionInstruction::CloneWithNewOperandsImpl(
@@ -1994,12 +2000,21 @@
HloDynamicSliceInstruction::HloDynamicSliceInstruction(
const Shape& shape, HloInstruction* operand, HloInstruction* start_indices,
absl::Span<const int64> slice_sizes)
- : HloInstruction(HloOpcode::kDynamicSlice, shape),
+ : HloDynamicIndexInstruction(HloOpcode::kDynamicSlice, shape),
dynamic_slice_sizes_(slice_sizes.begin(), slice_sizes.end()) {
AppendOperand(operand);
AppendOperand(start_indices);
}
+HloDynamicUpdateSliceInstruction::HloDynamicUpdateSliceInstruction(
+ const Shape& shape, HloInstruction* operand, HloInstruction* update,
+ HloInstruction* start_indices)
+ : HloDynamicIndexInstruction(HloOpcode::kDynamicUpdateSlice, shape) {
+ AppendOperand(operand);
+ AppendOperand(update);
+ AppendOperand(start_indices);
+}
+
HloInstructionProto HloDynamicSliceInstruction::ToProto() const {
HloInstructionProto proto = HloInstruction::ToProto();
for (int64 slice_size : dynamic_slice_sizes_) {
diff --git a/tensorflow/compiler/xla/service/hlo_instructions.h b/tensorflow/compiler/xla/service/hlo_instructions.h
index b5c2813..5420d4c 100644
--- a/tensorflow/compiler/xla/service/hlo_instructions.h
+++ b/tensorflow/compiler/xla/service/hlo_instructions.h
@@ -1171,7 +1171,14 @@
PaddingConfig padding_config_;
};
-class HloDynamicSliceInstruction : public HloInstruction {
+class HloDynamicIndexInstruction : public HloInstruction {
+ public:
+ explicit HloDynamicIndexInstruction(HloOpcode opcode, const Shape& shape)
+ : HloInstruction(opcode, shape) {}
+ virtual int64 index_operand_number() const = 0;
+};
+
+class HloDynamicSliceInstruction : public HloDynamicIndexInstruction {
public:
explicit HloDynamicSliceInstruction(const Shape& shape,
HloInstruction* operand,
@@ -1189,6 +1196,8 @@
// Returns a serialized representation of this instruction.
HloInstructionProto ToProto() const override;
+ int64 index_operand_number() const override { return 1; }
+
private:
std::vector<string> ExtraAttributesToStringImpl(
const HloPrintOptions& options) const override;
@@ -1206,6 +1215,16 @@
std::vector<int64> dynamic_slice_sizes_;
};
+class HloDynamicUpdateSliceInstruction : public HloDynamicIndexInstruction {
+ public:
+ explicit HloDynamicUpdateSliceInstruction(const Shape& shape,
+ HloInstruction* operand,
+ HloInstruction* update,
+ HloInstruction* start_indices);
+
+ int64 index_operand_number() const override { return 2; }
+};
+
class HloGatherInstruction : public HloInstruction {
public:
explicit HloGatherInstruction(
diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h
index 7b9cbf9..f1310e4 100644
--- a/tensorflow/compiler/xla/service/hlo_module.h
+++ b/tensorflow/compiler/xla/service/hlo_module.h
@@ -136,7 +136,9 @@
// information on opcode, shape, operands, and typically a root instruction.
// This function returns the same hash value for equivalent HLO modules,
// with respect to HloInstruction::Identical() method.
- uint64 Hash() const { return entry_computation()->Hash(); }
+ uint64 Hash() const {
+ return entry_computation()->root_instruction()->Hash();
+ }
// Gets the computations in this module.
//
diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
index 312b5d0..33ce7e2 100644
--- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
+++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
@@ -77,6 +77,11 @@
auto repeated_field = debug_options.xla_disable_hlo_passes();
absl::flat_hash_set<string> disabled_pass_names(repeated_field.begin(),
repeated_field.end());
+ if (debug_options.xla_disable_all_hlo_passes()) {
+ VLOG(1) << "*All* passes disabled by --xla_disable_all_hlo_passes.";
+ return {};
+ }
+
if (!disabled_pass_names.empty()) {
VLOG(1) << "Passes disabled by --xla_disable_hlo_passes: "
<< absl::StrJoin(disabled_pass_names, ", ");
@@ -113,7 +118,7 @@
}
const string message =
- StrCat("after ", after_pass_name, ", before ", before_pass_name);
+ absl::StrCat("after ", after_pass_name, ", before ", before_pass_name);
hlo_graph_dumper::MaybeDumpHloModule(module, message);
VLOG(3) << "HLO " << message << ":";
VLOG(3) << module.entry_computation_layout().ToString();
diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.cc b/tensorflow/compiler/xla/service/hlo_proto_util.cc
index 981d06c..3a9ee57 100644
--- a/tensorflow/compiler/xla/service/hlo_proto_util.cc
+++ b/tensorflow/compiler/xla/service/hlo_proto_util.cc
@@ -39,6 +39,7 @@
StatusOr<std::unique_ptr<HloModule>> CreateModuleFromProto(
const HloModuleProto& proto, const HloModuleConfig& module_config) {
+ VLOG(4) << proto.ShortDebugString();
TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
HloModule::CreateFromProto(proto, module_config));
TF_RETURN_IF_ERROR(
diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc
index 77db7b0..ace854e 100644
--- a/tensorflow/compiler/xla/service/hlo_verifier.cc
+++ b/tensorflow/compiler/xla/service/hlo_verifier.cc
@@ -481,7 +481,9 @@
const Shape& operand_shape_with_layout =
custom_call->operand_shapes_with_layout()[i];
TF_RET_CHECK(ShapeUtil::Compatible(custom_call->operand(i)->shape(),
- operand_shape_with_layout));
+ operand_shape_with_layout))
+ << custom_call->operand(i)->shape().ToString() << " operand "
+ << operand_shape_with_layout.ToString();
TF_RET_CHECK(LayoutUtil::HasLayout(operand_shape_with_layout));
}
}
diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc
index 3a5177c..d37ae94 100644
--- a/tensorflow/compiler/xla/service/interpreter/compiler.cc
+++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc
@@ -76,9 +76,12 @@
// need to compile anything
// Create executable from only the Hlo module.
+ auto evaluator = absl::make_unique<HloEvaluator>();
+ evaluator->set_use_fast_path(
+ hlo_module->config().debug_options().xla_hlo_evaluator_use_fast_path());
std::unique_ptr<Executable> executable =
- absl::make_unique<InterpreterExecutable>(
- std::move(hlo_module), absl::make_unique<HloEvaluator>());
+ absl::make_unique<InterpreterExecutable>(std::move(hlo_module),
+ std::move(evaluator));
return std::move(executable);
}
diff --git a/tensorflow/compiler/xla/service/interpreter/executable.cc b/tensorflow/compiler/xla/service/interpreter/executable.cc
index 7635fbf..de92040 100644
--- a/tensorflow/compiler/xla/service/interpreter/executable.cc
+++ b/tensorflow/compiler/xla/service/interpreter/executable.cc
@@ -85,6 +85,7 @@
Literal result_literal;
{
tensorflow::mutex_lock lock(evaluator_lock_);
+ evaluator_->ResetVisitStates();
TF_ASSIGN_OR_RETURN(result_literal, evaluator_->Evaluate<Literal>(
*computation, arg_literals));
}
diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc
index 311bd78..5c661bf 100644
--- a/tensorflow/compiler/xla/service/layout_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc
@@ -27,7 +27,6 @@
#include "tensorflow/compiler/xla/service/computation_layout.h"
#include "tensorflow/compiler/xla/service/hlo_computation.h"
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
-#include "tensorflow/compiler/xla/service/hlo_matchers.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_opcode.h"
#include "tensorflow/compiler/xla/service/hlo_parser.h"
diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc
index bd0139f..5eeb29c 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc
@@ -18,28 +18,29 @@
#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
namespace xla {
-Status KernelSupportLibrary::For(
+Status KernelSupportLibrary::ForWithStatus(
absl::string_view name, llvm::Value* start, llvm::Value* end,
llvm::Value* step,
const std::function<Status(llvm::Value*, bool)>& for_body_generator) {
- return If(b_->CreateICmpSLT(start, end), [&]() -> Status {
+ return IfWithStatus(b_->CreateICmpSLT(start, end), [&]() -> Status {
TF_RETURN_IF_ERROR(for_body_generator(start, /*is_first_iteration=*/true));
- return For(name, b_->CreateAdd(start, step), end, step,
- [&](llvm::Value* iv) { return for_body_generator(iv, false); });
+ return ForWithStatus(
+ name, b_->CreateAdd(start, step), end, step,
+ [&](llvm::Value* iv) { return for_body_generator(iv, false); });
});
}
-Status KernelSupportLibrary::For(
+Status KernelSupportLibrary::ForWithStatus(
absl::string_view name, llvm::Value* start, llvm::Value* end,
llvm::Value* step, bool peel_first_iteration,
const std::function<Status(llvm::Value*, llvm::Value*)>&
for_body_generator) {
if (peel_first_iteration) {
- return For(name, start, end, step, true,
- [&](llvm::Value* indvar, bool is_first_iteration) -> Status {
- return for_body_generator(indvar,
- b_->getInt1(is_first_iteration));
- });
+ return ForWithStatus(
+ name, start, end, step, true,
+ [&](llvm::Value* indvar, bool is_first_iteration) -> Status {
+ return for_body_generator(indvar, b_->getInt1(is_first_iteration));
+ });
} else {
std::unique_ptr<llvm_ir::ForLoop> loop = llvm_ir::ForLoop::EmitForLoop(
name, start, end, step, b_,
@@ -55,7 +56,7 @@
}
}
-Status KernelSupportLibrary::If(
+Status KernelSupportLibrary::IfWithStatus(
absl::string_view name, llvm::Value* condition,
const std::function<Status()>& true_block_generator,
const std::function<Status()>& false_block_generator) {
diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h
index 43fec31..612b839 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h
+++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h
@@ -48,41 +48,42 @@
// for (i64 i = `start` + `step`; i s< `end`; i += `step`)
// `for_body_generator(/*ind_var=*/,i, /*is_first_iteration=*/false)`;
// }
- Status For(
+ Status ForWithStatus(
absl::string_view name, llvm::Value* start, llvm::Value* end,
llvm::Value* step,
const std::function<Status(llvm::Value* ind_var,
bool is_first_iteration)>& for_body_generator);
- void ForReturnVoid(
+ void For(
absl::string_view name, llvm::Value* start, llvm::Value* end,
llvm::Value* step,
const std::function<void(llvm::Value* ind_var, bool is_first_iteration)>&
for_body_generator) {
CHECK_EQ(Status::OK(),
- For(name, start, end, step,
+ ForWithStatus(
+ name, start, end, step,
[&](llvm::Value* ind_var, bool is_first_iteration) -> Status {
for_body_generator(ind_var, is_first_iteration);
return Status::OK();
}));
}
- Status For(absl::string_view name, int64 start, int64 end, int64 step,
- const std::function<Status(llvm::Value* ind_var,
- bool is_first_iteration)>&
- for_body_generator) {
- return For(name, /*start=*/b_->getInt64(start),
- /*end=*/b_->getInt64(end),
- /*step=*/b_->getInt64(step), for_body_generator);
+ Status ForWithStatus(
+ absl::string_view name, int64 start, int64 end, int64 step,
+ const std::function<Status(
+ llvm::Value* ind_var, bool is_first_iteration)>& for_body_generator) {
+ return ForWithStatus(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
- void ForReturnVoid(
+ void For(
absl::string_view name, int64 start, int64 end, int64 step,
const std::function<void(llvm::Value* ind_var, bool is_first_iteration)>&
for_body_generator) {
- ForReturnVoid(name, /*start=*/b_->getInt64(start),
- /*end=*/b_->getInt64(end),
- /*step=*/b_->getInt64(step), for_body_generator);
+ For(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
// Generates the following control flow structure if `peel_first_iteration` is
@@ -99,19 +100,19 @@
// for (i64 i = `start`; i s< `end`; i += `step`)
// `for_body_generator(/*ind_var=*/,i,
// /*is_first_iteration=*/,(i != `start`))`;
- Status For(absl::string_view name, llvm::Value* start, llvm::Value* end,
- llvm::Value* step, bool peel_first_iteration,
- const std::function<Status(llvm::Value* ind_var,
- llvm::Value* is_first_iteration)>&
- for_body_generator);
+ Status ForWithStatus(
+ absl::string_view name, llvm::Value* start, llvm::Value* end,
+ llvm::Value* step, bool peel_first_iteration,
+ const std::function<Status(llvm::Value* ind_var,
+ llvm::Value* is_first_iteration)>&
+ for_body_generator);
- void ForReturnVoid(absl::string_view name, llvm::Value* start,
- llvm::Value* end, llvm::Value* step,
- bool peel_first_iteration,
- const std::function<void(llvm::Value* ind_var,
- llvm::Value* is_first_iteration)>&
- for_body_generator) {
- TF_CHECK_OK(For(
+ void For(absl::string_view name, llvm::Value* start, llvm::Value* end,
+ llvm::Value* step, bool peel_first_iteration,
+ const std::function<void(llvm::Value* ind_var,
+ llvm::Value* is_first_iteration)>&
+ for_body_generator) {
+ TF_CHECK_OK(ForWithStatus(
name, start, end, step, peel_first_iteration,
[&](llvm::Value* ind_var, llvm::Value* is_first_iteration) -> Status {
for_body_generator(ind_var, is_first_iteration);
@@ -119,80 +120,81 @@
}));
}
- Status For(absl::string_view name, llvm::Value* start, llvm::Value* end,
- int64 step, bool peel_first_iteration,
- const std::function<Status(llvm::Value* ind_var,
- llvm::Value* is_first_iteration)>&
- for_body_generator) {
- return For(name, /*start=*/start, /*end=*/end,
- /*step=*/llvm::ConstantInt::get(start->getType(), step),
- peel_first_iteration, for_body_generator);
+ Status ForWithStatus(
+ absl::string_view name, llvm::Value* start, llvm::Value* end, int64 step,
+ bool peel_first_iteration,
+ const std::function<Status(llvm::Value* ind_var,
+ llvm::Value* is_first_iteration)>&
+ for_body_generator) {
+ return ForWithStatus(
+ name, /*start=*/start, /*end=*/end,
+ /*step=*/llvm::ConstantInt::get(start->getType(), step),
+ peel_first_iteration, for_body_generator);
}
- void ForReturnVoid(absl::string_view name, llvm::Value* start,
- llvm::Value* end, int64 step, bool peel_first_iteration,
- const std::function<void(llvm::Value* ind_var,
- llvm::Value* is_first_iteration)>&
- for_body_generator) {
- ForReturnVoid(name, /*start=*/start, /*end=*/end,
- /*step=*/llvm::ConstantInt::get(start->getType(), step),
- peel_first_iteration, for_body_generator);
+ void For(absl::string_view name, llvm::Value* start, llvm::Value* end,
+ int64 step, bool peel_first_iteration,
+ const std::function<void(llvm::Value* ind_var,
+ llvm::Value* is_first_iteration)>&
+ for_body_generator) {
+ For(name, /*start=*/start, /*end=*/end,
+ /*step=*/llvm::ConstantInt::get(start->getType(), step),
+ peel_first_iteration, for_body_generator);
}
- Status For(
+ Status ForWithStatus(
absl::string_view name, llvm::Value* start, llvm::Value* end,
llvm::Value* step,
const std::function<Status(llvm::Value* ind_var)>& for_body_generator) {
- return For(name, start, end, step,
- /*peel_first_iteration=*/false,
- [&](llvm::Value* indvar, llvm::Value*) -> Status {
- return for_body_generator(indvar);
- });
+ return ForWithStatus(name, start, end, step,
+ /*peel_first_iteration=*/false,
+ [&](llvm::Value* indvar, llvm::Value*) -> Status {
+ return for_body_generator(indvar);
+ });
}
- void ForReturnVoid(
+ void For(
absl::string_view name, llvm::Value* start, llvm::Value* end,
llvm::Value* step,
const std::function<void(llvm::Value* ind_var)>& for_body_generator) {
- ForReturnVoid(name, start, end, step,
- /*peel_first_iteration=*/false,
- [&](llvm::Value* indvar, llvm::Value*) {
- return for_body_generator(indvar);
- });
+ For(name, start, end, step,
+ /*peel_first_iteration=*/false, [&](llvm::Value* indvar, llvm::Value*) {
+ return for_body_generator(indvar);
+ });
}
- Status For(
+ Status ForWithStatus(
absl::string_view name, llvm::Value* start, llvm::Value* end, int64 step,
const std::function<Status(llvm::Value* ind_var)>& for_body_generator) {
- return For(name, start, end, llvm::ConstantInt::get(start->getType(), step),
- /*peel_first_iteration=*/false,
- [&](llvm::Value* indvar, llvm::Value*) -> Status {
- return for_body_generator(indvar);
- });
+ return ForWithStatus(name, start, end,
+ llvm::ConstantInt::get(start->getType(), step),
+ /*peel_first_iteration=*/false,
+ [&](llvm::Value* indvar, llvm::Value*) -> Status {
+ return for_body_generator(indvar);
+ });
}
- void ForReturnVoid(
+ void For(
absl::string_view name, llvm::Value* start, llvm::Value* end, int64 step,
const std::function<void(llvm::Value* ind_var)>& for_body_generator) {
- ForReturnVoid(name, start, end,
- llvm::ConstantInt::get(start->getType(), step),
- for_body_generator);
+ For(name, start, end, llvm::ConstantInt::get(start->getType(), step),
+ for_body_generator);
}
- Status For(
+ Status ForWithStatus(
absl::string_view name, int64 start, int64 end, int64 step,
const std::function<Status(llvm::Value* ind_var)>& for_body_generator) {
- return For(name, /*start=*/b_->getInt64(start),
- /*end=*/b_->getInt64(end),
- /*step=*/b_->getInt64(step), for_body_generator);
+ return ForWithStatus(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
- void ForReturnVoid(
+ void For(
absl::string_view name, int64 start, int64 end, int64 step,
const std::function<void(llvm::Value* ind_var)>& for_body_generator) {
- ForReturnVoid(name, /*start=*/b_->getInt64(start),
- /*end=*/b_->getInt64(end),
- /*step=*/b_->getInt64(step), for_body_generator);
+ For(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
// Generates the following control flow structure:
@@ -201,38 +203,43 @@
// `true_block_generator()`;
// else
// `false_block_generator()`;
- Status If(absl::string_view name, llvm::Value* condition,
- const std::function<Status()>& true_block_generator,
- const std::function<Status()>& false_block_generator =
- []() -> Status { return Status::OK(); });
+ Status IfWithStatus(
+ absl::string_view name, llvm::Value* condition,
+ const std::function<Status()>& true_block_generator,
+ const std::function<Status()>& false_block_generator = []() -> Status {
+ return Status::OK();
+ });
- Status If(llvm::Value* condition,
- const std::function<Status()>& true_block_generator,
- const std::function<Status()>& false_block_generator =
- []() -> Status { return Status::OK(); }) {
- return If("", condition, true_block_generator, false_block_generator);
+ Status IfWithStatus(
+ llvm::Value* condition,
+ const std::function<Status()>& true_block_generator,
+ const std::function<Status()>& false_block_generator = []() -> Status {
+ return Status::OK();
+ }) {
+ return IfWithStatus("", condition, true_block_generator,
+ false_block_generator);
}
- void IfReturnVoid(llvm::Value* condition,
- const std::function<void()>& true_block_generator,
- const std::function<void()>& false_block_generator = []() {
- }) {
- IfReturnVoid("", condition, true_block_generator, false_block_generator);
+ void If(
+ llvm::Value* condition, const std::function<void()>& true_block_generator,
+ const std::function<void()>& false_block_generator = []() {}) {
+ If("", condition, true_block_generator, false_block_generator);
}
- void IfReturnVoid(absl::string_view name, llvm::Value* condition,
- const std::function<void()>& true_block_generator,
- const std::function<void()>& false_block_generator = []() {
- }) {
- TF_CHECK_OK(If(name, condition,
- [&]() {
- true_block_generator();
- return Status::OK();
- },
- [&]() {
- false_block_generator();
- return Status::OK();
- }));
+ void If(
+ absl::string_view name, llvm::Value* condition,
+ const std::function<void()>& true_block_generator,
+ const std::function<void()>& false_block_generator = []() {}) {
+ TF_CHECK_OK(IfWithStatus(
+ name, condition,
+ [&]() {
+ true_block_generator();
+ return Status::OK();
+ },
+ [&]() {
+ false_block_generator();
+ return Status::OK();
+ }));
}
using ArgumentVector = absl::Span<llvm::Value* const>;
diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc
index 1aa85eb..cebbc42 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc
@@ -170,14 +170,16 @@
IrArray::Index KernelMappingScheme::GetTileIndexForBlockOrigin(
const IrArray::Index& block_index) {
- IrArray::Index tile_index = block_index;
+ DCHECK_EQ(block_index.size(), block_sizes_.size());
+ std::vector<llvm::Value*> multidim;
+ multidim.reserve(block_sizes_.size());
for (int i = 0; i < block_sizes_.size(); ++i) {
- tile_index[i] = b_->CreateMul(
+ multidim.push_back(b_->CreateMul(
block_index[i],
llvm::ConstantInt::get(block_index[i]->getType(), block_sizes_[i]),
- "block_origin." + std::to_string(i));
+ "block_origin." + std::to_string(i)));
}
- return tile_index;
+ return IrArray::Index(multidim, block_index[0]->getType());
}
IrArray::Index KernelMappingScheme::GetElementIndexForTileOrigin(
@@ -217,14 +219,14 @@
// defined by (num_thread_y, num_thread_x) from thread_id.
llvm::CallInst* thread_id_raw = llvm_ir::EmitCallToIntrinsic(
llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}, b_);
- llvm_ir::AddRangeMetadata(0, GetThreadsPerTile(), thread_id_raw);
+ llvm_ir::AddRangeMetadata(0, GetThreadsPerBlock(), thread_id_raw);
llvm::Value* thread_id_int =
b_->CreateIntCast(thread_id_raw, index_ty,
/*isSigned=*/true, "thread.id.x");
llvm::Value* num_thread_x =
llvm::ConstantInt::get(index_ty, GetNumberOfThreadsForDimensionX());
- llvm::Value* x = b_->CreateURem(thread_id_int, num_thread_x);
- llvm::Value* y = b_->CreateUDiv(thread_id_int, num_thread_x);
+ llvm::Value* x = b_->CreateURem(thread_id_int, num_thread_x, "thread.x");
+ llvm::Value* y = b_->CreateUDiv(thread_id_int, num_thread_x, "thread.y");
return std::make_tuple(y, x);
}
diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h
index 7277aea..fb633b1 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h
+++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h
@@ -142,7 +142,7 @@
int64 GetNumberOfThreadsForDimensionX() const { return num_threads_x_; }
int64 GetNumberOfThreadsForDimensionY() const { return num_threads_y_; }
- int64 GetThreadsPerTile() const {
+ int64 GetThreadsPerBlock() const {
return GetNumberOfThreadsForDimensionX() *
GetNumberOfThreadsForDimensionY();
}
diff --git a/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc b/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc
index e22c217..6a9406b 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc
@@ -108,7 +108,7 @@
// if (is_smaller_index && index_is_inbounds)
KernelSupportLibrary ksl(b);
- ksl.IfReturnVoid("smaller_comparison_index", do_comparison, [&]() {
+ ksl.If("smaller_comparison_index", do_comparison, [&]() {
auto key1 = read_element(0, current_keys_index);
auto key2 = read_element(0, compare_keys_index);
auto compare_key1 = key1;
@@ -155,7 +155,7 @@
is_smaller_than = b->CreateOr(
is_smaller_than, b->CreateAnd(keys_equal, index_is_smaller_than));
}
- ksl.IfReturnVoid("is_smaller_than", is_smaller_than, [&]() {
+ ksl.If("is_smaller_than", is_smaller_than, [&]() {
// Swap key1 with key2.
write_element(0, current_keys_index, key2);
write_element(0, compare_keys_index, key1);
@@ -192,7 +192,7 @@
b->CreateShl(tiled_keys_index[dimension_to_sort], value_one);
// We want to copy two adjacent elements. We first check whether the
// first index position is within bounds.
- ksl.IfReturnVoid(
+ ksl.If(
"smaller_keys_index",
b->CreateICmpSLT(current_keys_index,
tiled_keys_index.GetConstantWithIndexType(
@@ -203,15 +203,14 @@
// Increment to go the next index position.
current_keys_index = b->CreateAdd(current_keys_index, value_one);
// Here we check whether the next index position is within bounds.
- ksl.IfReturnVoid(
- "inner_smaller_keys_index",
- b->CreateICmpSLT(current_keys_index,
- tiled_keys_index.GetConstantWithIndexType(
- dimension_to_sort_bound)),
- [&]() {
- cache_index = b->CreateAdd(cache_index, value_one);
- read_or_write(cache_index, current_keys_index);
- });
+ ksl.If("inner_smaller_keys_index",
+ b->CreateICmpSLT(current_keys_index,
+ tiled_keys_index.GetConstantWithIndexType(
+ dimension_to_sort_bound)),
+ [&]() {
+ cache_index = b->CreateAdd(cache_index, value_one);
+ read_or_write(cache_index, current_keys_index);
+ });
});
};
@@ -253,7 +252,7 @@
if (dimension_to_sort_bound % tile_size) {
// Otherwise we need a bounds check for the last tile. The last tile has
// size 'dimension_to_sort_bound' % 'tile_size'.
- ksl.IfReturnVoid(
+ ksl.If(
"is_last_tile",
b->CreateICmpUGE(
b->CreateMul(tiled_keys_index[dimension_to_sort],
diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h
index 432aa1e..c35f726 100644
--- a/tensorflow/compiler/xla/service/pattern_matcher.h
+++ b/tensorflow/compiler/xla/service/pattern_matcher.h
@@ -2232,6 +2232,7 @@
// We could implement all ops as "variadic" ops, but it would make the
// already-bad compile errors even worse.
+XLA_VARIADIC_OP_PATTERN(AfterAll);
XLA_VARIADIC_OP_PATTERN(Concatenate);
XLA_VARIADIC_OP_PATTERN(CustomCall);
XLA_VARIADIC_OP_PATTERN(Map)
diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc
index 5ec7fe2..ae5bd93 100644
--- a/tensorflow/compiler/xla/service/service.cc
+++ b/tensorflow/compiler/xla/service/service.cc
@@ -1078,9 +1078,11 @@
ProgramShape program_shape(arg->computation().host_program_shape());
TF_DCHECK_OK(ShapeUtil::ValidateShape(program_shape.result()));
+ absl::optional<Layout> output_layout;
if (arg->has_output_layout()) {
+ output_layout = Layout::CreateFromProto(arg->output_layout());
TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutForShape(
- arg->output_layout(), program_shape.result()));
+ *output_layout, program_shape.result()));
}
HloModuleConfig config(program_shape);
@@ -1096,8 +1098,8 @@
// relayout here.
//
// TODO(b/77824332): Make HloEvaluator take care of the re-layout.
- if (arg->has_output_layout()) {
- result_literal = result_literal.Relayout(arg->output_layout());
+ if (output_layout.has_value()) {
+ result_literal = result_literal.Relayout(*output_layout);
}
*result->mutable_literal() = result_literal.ToProto();
diff --git a/tensorflow/compiler/xla/shape.cc b/tensorflow/compiler/xla/shape.cc
index 746ab9e..b206345 100644
--- a/tensorflow/compiler/xla/shape.cc
+++ b/tensorflow/compiler/xla/shape.cc
@@ -32,7 +32,7 @@
*add_tuple_shapes() = Shape(element_shape);
}
if (shape_proto.has_layout()) {
- *mutable_layout() = shape_proto.layout();
+ *mutable_layout() = Layout::CreateFromProto(shape_proto.layout());
}
}
@@ -48,7 +48,7 @@
*proto.add_tuple_shapes() = shape.ToProto();
}
if (has_layout()) {
- *proto.mutable_layout() = layout();
+ *proto.mutable_layout() = layout().ToProto();
}
return proto;
}
diff --git a/tensorflow/compiler/xla/shape.h b/tensorflow/compiler/xla/shape.h
index 7f6b14a..7643f64 100644
--- a/tensorflow/compiler/xla/shape.h
+++ b/tensorflow/compiler/xla/shape.h
@@ -20,6 +20,7 @@
#include <vector>
#include "absl/types/optional.h"
+#include "tensorflow/compiler/xla/layout.h"
#include "tensorflow/compiler/xla/types.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/platform/types.h"
@@ -76,21 +77,10 @@
std::vector<Shape>* mutable_tuple_shapes() { return &tuple_shapes_; }
// Methods for accessing the layout field.
- bool has_layout() const { return layout_.has_value(); }
- const Layout& layout() const {
- if (layout_.has_value()) {
- return *layout_;
- } else {
- return Layout::default_instance();
- }
- }
- Layout* mutable_layout() {
- if (!layout_.has_value()) {
- layout_ = Layout();
- }
- return &layout_.value();
- }
- void clear_layout() { layout_.reset(); }
+ bool has_layout() const { return layout_.format() != INVALID_FORMAT; }
+ const Layout& layout() const { return layout_; }
+ Layout* mutable_layout() { return &layout_; }
+ void clear_layout() { layout_.Clear(); }
void Swap(Shape* other) {
using std::swap;
@@ -101,7 +91,7 @@
element_type_ = PRIMITIVE_TYPE_INVALID;
dimensions_.clear();
tuple_shapes_.clear();
- layout_.reset();
+ clear_layout();
}
string SerializeAsString() const { return ToProto().SerializeAsString(); }
@@ -118,8 +108,8 @@
// The tuple element subshapes. This is nonempty only for tuple shapes.
std::vector<Shape> tuple_shapes_;
- // The array layout of the shape. This is present only for array shapes.
- absl::optional<Layout> layout_;
+ // The layout of the shape. Only relevant for arrays.
+ Layout layout_;
};
// Shape of the parameters and output of an XLA computation. This is analogous
diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc
index f3cc51c..eef2dc9 100644
--- a/tensorflow/compiler/xla/shape_util.cc
+++ b/tensorflow/compiler/xla/shape_util.cc
@@ -164,9 +164,9 @@
TF_ASSIGN_OR_RETURN(Shape shape,
ShapeUtil::MakeValidatedShape(element_type, dimensions));
auto min2maj = shape.mutable_layout()->mutable_minor_to_major();
- min2maj->Clear();
+ min2maj->clear();
for (int64 value : minor_to_major) {
- min2maj->Add(value);
+ min2maj->push_back(value);
}
if (!shape.has_layout()) {
return InvalidArgument("Shape has no layout.");
@@ -584,7 +584,7 @@
// Parses shapes with simple recursive descent structure -- consumes from the
// front of s and passes that view recursively as required.
StatusOr<Shape> ParseShapeStringInternal(absl::string_view* s) {
- *s = StripLeadingAsciiWhitespace(*s);
+ *s = absl::StripLeadingAsciiWhitespace(*s);
if (absl::ConsumePrefix(s, "(")) { // Tuple.
std::vector<Shape> shapes;
@@ -597,7 +597,7 @@
}
shapes.emplace_back();
TF_ASSIGN_OR_RETURN(shapes.back(), ParseShapeStringInternal(s));
- *s = StripLeadingAsciiWhitespace(*s);
+ *s = absl::StripLeadingAsciiWhitespace(*s);
must_end = !absl::ConsumePrefix(s, ",");
}
return ShapeUtil::MakeTupleShape(shapes);
@@ -1618,10 +1618,10 @@
if (LayoutUtil::HasLayout(shape)) {
Layout* layout = shape.mutable_layout();
layout->set_format(DENSE);
- for (size_t i = 0; i < layout->minor_to_major().size();) {
+ for (int64 i = 0; i < layout->minor_to_major().size();) {
if (layout->minor_to_major(i) == dim_to_delete) {
layout->mutable_minor_to_major()->erase(
- layout->minor_to_major().begin() + i);
+ layout->mutable_minor_to_major()->begin() + i);
continue;
}
if (layout->minor_to_major(i) > dim_to_delete) {
diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index f7f090f..0300b64 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -1,6 +1,13 @@
# Description:
# Base testing infrastructure for XLA.
+load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites", "generate_backend_test_macros", "xla_test", "xla_test_library")
+load(
+ "//tensorflow/core:platform/default/build_config_root.bzl",
+ "tf_cuda_tests_tags",
+)
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
+
licenses(["notice"]) # Apache 2.0
package(
@@ -23,17 +30,6 @@
]),
)
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test")
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test_library")
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites")
-load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_test_macros")
-load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
-load("//tensorflow:tensorflow.bzl", "tf_cc_test")
-load(
- "//tensorflow/core:platform/default/build_config_root.bzl",
- "tf_cuda_tests_tags",
-)
-
# Generate test_suites for all backends, named "${backend}_tests".
generate_backend_suites()
@@ -303,10 +299,6 @@
name = "conv_depthwise_test",
timeout = "long",
srcs = ["conv_depthwise_test.cc"],
- blacklisted_backends = [
- # disabled because of a break b/119590850.
- "gpu",
- ],
shard_count = 50,
deps = [
"//tensorflow/compiler/xla:execution_options_util",
@@ -1352,6 +1344,7 @@
xla_test(
name = "custom_call_test",
srcs = ["custom_call_test.cc"],
+ backends = ["cpu"],
deps = [
"//tensorflow/compiler/xla:literal",
"//tensorflow/compiler/xla:literal_util",
diff --git a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc
index f6be27b..915b456 100644
--- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc
+++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc
@@ -329,13 +329,13 @@
Literal b_literal = LiteralUtil::CreateR1<float>({b_values});
std::unique_ptr<GlobalData> b_data =
client_->TransferToServer(b_literal).ConsumeValueOrDie();
- auto b_constant = Parameter(&builder, 1, a_literal.shape(), "b_param");
- auto b_param = ConstantR1<float>(&builder, b_values);
+ auto b_param = Parameter(&builder, 1, a_literal.shape(), "b_param");
+ auto b_constant = ConstantR1<float>(&builder, b_values);
- auto sum1 = Add(a_constant, b_constant);
- auto sum2 = Add(a_constant, b_param);
- auto sum3 = Add(a_param, b_constant);
- auto sum4 = Add(a_param, b_param);
+ auto sum1 = Add(a_constant, b_param);
+ auto sum2 = Add(a_constant, b_constant);
+ auto sum3 = Add(a_param, b_param);
+ auto sum4 = Add(a_param, b_constant);
auto sum = Add(sum1, sum2);
sum = Add(sum, sum3);
diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc
index 12c0299..697236d 100644
--- a/tensorflow/compiler/xla/tests/client_library_test_base.cc
+++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc
@@ -74,6 +74,9 @@
// default.
execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes(
"constant_folding");
+
+ execution_options_.mutable_debug_options()
+ ->set_xla_hlo_evaluator_use_fast_path(true);
}
ClientLibraryTestBase::ClientLibraryTestBase(se::Platform* platform)
@@ -88,6 +91,9 @@
execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes(
"constant_folding");
+
+ execution_options_.mutable_debug_options()
+ ->set_xla_hlo_evaluator_use_fast_path(true);
}
string ClientLibraryTestBase::TestName() const {
diff --git a/tensorflow/compiler/xla/tests/conv_depthwise_test.cc b/tensorflow/compiler/xla/tests/conv_depthwise_test.cc
index bc9bd8a..627a17a 100644
--- a/tensorflow/compiler/xla/tests/conv_depthwise_test.cc
+++ b/tensorflow/compiler/xla/tests/conv_depthwise_test.cc
@@ -50,9 +50,9 @@
static std::vector<DepthwiseConvolution2DSpec> GetConv2DTestCases() {
std::vector<DepthwiseConvolution2DSpec> config_set;
std::vector<std::vector<int64>> config_options = {
- {128, 6, 3, 64}, {256, 5, 3, 256}, {256, 5, 2, 144}, {144, 5, 3, 64},
- {144, 5, 2, 256}, {8, 48, 17, 8}, {128, 20, 6, 64}, {128, 1, 2, 144},
- {256, 1, 2, 64}, {64, 14, 12, 172}, {16, 9, 4, 16}};
+ {128, 6, 3, 64}, {256, 5, 3, 256}, {256, 5, 2, 144}, {144, 5, 3, 64},
+ {144, 5, 2, 256}, {8, 48, 17, 8}, {128, 20, 6, 64}, {64, 14, 12, 172},
+ {16, 9, 4, 16}, {128, 1, 2, 144}, {256, 1, 2, 64}};
for (auto option : config_options) {
int64 feature = option[0];
diff --git a/tensorflow/compiler/xla/tests/copy_test.cc b/tensorflow/compiler/xla/tests/copy_test.cc
index 3622f2c..df005a6 100644
--- a/tensorflow/compiler/xla/tests/copy_test.cc
+++ b/tensorflow/compiler/xla/tests/copy_test.cc
@@ -133,7 +133,9 @@
// Reverse the minor-to-major order of the literal.
Layout* literal_layout = literal.mutable_shape_do_not_use()->mutable_layout();
ASSERT_EQ(2, literal_layout->minor_to_major_size());
- literal_layout->mutable_minor_to_major()->SwapElements(0, 1);
+ // Swap the first and second elements.
+ *literal_layout->mutable_minor_to_major() = {
+ literal_layout->minor_to_major(1), literal_layout->minor_to_major(0)};
HloInstruction* constant = builder.AddInstruction(
HloInstruction::CreateConstant(std::move(literal)));
diff --git a/tensorflow/compiler/xla/tests/custom_call_test.cc b/tensorflow/compiler/xla/tests/custom_call_test.cc
index 738b644..cad43d1 100644
--- a/tensorflow/compiler/xla/tests/custom_call_test.cc
+++ b/tensorflow/compiler/xla/tests/custom_call_test.cc
@@ -54,11 +54,20 @@
out[2] = array[2] + 1;
out[3] = array[3] + 1;
}
+
+void F32TupleSwap(float** out, float** in) {
+ TF_ANNOTATE_MEMORY_IS_INITIALIZED(in[0], sizeof(float));
+ TF_ANNOTATE_MEMORY_IS_INITIALIZED(in[1], sizeof(float));
+ *out[0] = *in[1];
+ *out[1] = *in[0];
+}
+
} // namespace
REGISTER_CUSTOM_CALL_TARGET(R0F32Add2);
REGISTER_CUSTOM_CALL_TARGET(R2F32ReduceSum);
REGISTER_CUSTOM_CALL_TARGET(Add1ToValues);
+REGISTER_CUSTOM_CALL_TARGET(F32TupleSwap);
namespace xla {
namespace {
@@ -69,7 +78,7 @@
Shape r2f32_ = ShapeUtil::MakeShape(F32, {2, 2});
};
-XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR0F32Add2)) {
+XLA_TEST_F(CustomCallTest, CustomCallR0F32Add2) {
auto module = CreateNewUnverifiedModule();
auto builder = HloComputation::Builder(TestName());
@@ -84,7 +93,7 @@
LiteralTestUtil::ExpectR0Near<float>(44.0f, result, error_spec_);
}
-XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR2F32Reduce)) {
+XLA_TEST_F(CustomCallTest, CustomCallR2F32Reduce) {
auto module = CreateNewUnverifiedModule();
auto builder = HloComputation::Builder(TestName());
@@ -105,7 +114,7 @@
LiteralTestUtil::ExpectR0Near<float>(10.0f, result, error_spec_);
}
-XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(UsedInOtherComputations)) {
+XLA_TEST_F(CustomCallTest, UsedInOtherComputations) {
auto module = CreateNewUnverifiedModule();
auto b = HloComputation::Builder(TestName());
@@ -129,7 +138,7 @@
Array3D<float>{{{2, 3}, {4, 5}}, {{3, 4}, {5, 6}}}, result);
}
-XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(InputAndOutputLayoutDiffer)) {
+XLA_TEST_F(CustomCallTest, InputAndOutputLayoutDiffer) {
auto module = CreateNewUnverifiedModule();
auto b = HloComputation::Builder(TestName());
@@ -151,7 +160,7 @@
LiteralTestUtil::ExpectR2Equal<float>({{2.f, 4.f}, {3.f, 5.f}}, result);
}
-XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(LayoutConstrained)) {
+XLA_TEST_F(CustomCallTest, LayoutConstrained) {
// The argument and result of the computation are set to different layouts,
// but the custom call is layout constrained to a fixed operand and result
// layout, so the correct result should be produced.
@@ -176,6 +185,26 @@
LiteralTestUtil::ExpectR2Equal<float>({{2.f, 3.f}, {4.f, 5.f}}, result);
}
+XLA_TEST_F(CustomCallTest, TupleOutput) {
+ const char* kModuleStr = R"(
+ HloModule m
+ test {
+ p0 = f32[] parameter(0)
+ p1 = f32[] parameter(1)
+ ROOT %custom-call = (f32[], f32[]) custom-call(f32[] %p0, f32[] %p1), custom_call_target="F32TupleSwap", operand_layout_constraints={f32[], f32[]}
+ }
+ )";
+ TF_ASSERT_OK_AND_ASSIGN(auto module,
+ ParseAndReturnVerifiedModule(kModuleStr));
+
+ Literal arg0 = LiteralUtil::CreateR0<float>(7.f);
+ Literal arg1 = LiteralUtil::CreateR0<float>(42.f);
+
+ Literal expected = LiteralUtil::MakeTuple({&arg1, &arg0});
+ Literal result = ExecuteAndTransfer(std::move(module), {&arg0, &arg1});
+ EXPECT_EQ(result, expected);
+}
+
class CustomCallClientAPITest : public ClientLibraryTestBase {};
// When using the client API, CustomCall targets can't begin with '$' -- these
diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc
index 25091b8..c5d8b66 100644
--- a/tensorflow/compiler/xla/tests/dot_operation_test.cc
+++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc
@@ -30,7 +30,6 @@
#include "tensorflow/compiler/xla/tests/test_utils.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
-#include "tensorflow/core/util/command_line_flags.h"
namespace xla {
namespace {
diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc
index 989a7c7..d57846e 100644
--- a/tensorflow/compiler/xla/tests/hlo_test_base.cc
+++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc
@@ -181,6 +181,7 @@
// TODO(b/38354253): Change tests to use Parameters instead of Constants.
debug_options.add_xla_disable_hlo_passes("constant_folding");
debug_options.set_xla_gpu_max_kernel_unroll_factor(1);
+ debug_options.set_xla_hlo_evaluator_use_fast_path(true);
return debug_options;
}
diff --git a/tensorflow/compiler/xla/tests/iota_test.cc b/tensorflow/compiler/xla/tests/iota_test.cc
index 65205f5..37b2c63 100644
--- a/tensorflow/compiler/xla/tests/iota_test.cc
+++ b/tensorflow/compiler/xla/tests/iota_test.cc
@@ -80,7 +80,7 @@
}
INSTANTIATE_TEST_CASE_P(IotaR2TestInstantiation, IotaR2Test,
- ::testing::Combine(::testing::Values(F32, S32),
+ ::testing::Combine(::testing::Values(F32, S32, BF16),
::testing::Range(/*start=*/10,
/*end=*/1001,
/*step=*/10),
diff --git a/tensorflow/compiler/xla/tools/replay_computation.cc b/tensorflow/compiler/xla/tools/replay_computation.cc
index ff2c339..1a51303 100644
--- a/tensorflow/compiler/xla/tools/replay_computation.cc
+++ b/tensorflow/compiler/xla/tools/replay_computation.cc
@@ -118,7 +118,12 @@
std::vector<std::unique_ptr<GlobalData>> global_data_arguments;
std::vector<const ShapedBuffer*> argument_ptrs;
if (opts.use_fake_data) {
- global_data_arguments = MakeFakeArgumentsOrDie(computation, client);
+ // Run fake computations with debug options ignoring XLA_FLAGS. Users very
+ // likely want XLA_FLAGS only to apply to the "real" computation being run,
+ // not to the fake computations we use for generating arguments.
+ auto debug_opts = DefaultDebugOptionsIgnoringFlags();
+ global_data_arguments =
+ MakeFakeArgumentsOrDie(computation, client, &debug_opts);
for (const auto& data : global_data_arguments) {
argument_ptrs.push_back(
client->GlobalDataToShapedBuffer(data->handle(), /*device_ordinal=*/0)
diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto
index bdeb172..0e8fa73 100644
--- a/tensorflow/compiler/xla/xla.proto
+++ b/tensorflow/compiler/xla/xla.proto
@@ -100,6 +100,14 @@
// names as specified by the HloPassInterface::name() method.
repeated string xla_disable_hlo_passes = 30;
+ // Disables all HLO passes. Notes that some passes are necessary for
+ // correctness and the invariants that must be satisfied by "fully optimized"
+ // HLO are different for different devices and may change over time. The only
+ // "guarantee", such as it is, is that if you compile XLA and dump the
+ // optimized HLO for some graph, you should be able to run it again on the
+ // same device with the same build of XLA.
+ bool xla_disable_all_hlo_passes = 104;
+
// Numerical optimization level for the XLA compiler backend; the specific
// interpretation of this value is left to the backends.
int32 xla_backend_optimization_level = 31;
@@ -213,6 +221,17 @@
// the host that run models in parallel across multiple devices.
int32 xla_force_host_platform_device_count = 102;
+ // If set to true XLA:GPU invokes `ptxas` with -O0 (default is -O3).
+ bool xla_gpu_disable_ptxas_optimizations = 103;
+
+ // Dump HLO graphs as an HTML (DOT -> SVG inlined in HTML)
+ bool xla_hlo_dump_as_html = 105;
+
+ // Enable fast math with eigen in the HLO evaluator.
+ bool xla_hlo_evaluator_use_fast_path = 106;
+
+ // Next id: 107
+
// Extra options to pass to the compilation backend (e.g. LLVM); specific
// interpretation of these values is left to the backend.
map<string, string> xla_backend_extra_options = 500;
@@ -386,7 +405,7 @@
message ComputeConstantGraphRequest {
HloModuleProto computation = 1;
- Layout output_layout = 2;
+ LayoutProto output_layout = 2;
}
message ComputeConstantResponse {
diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto
index 85ec834..e9c86ab 100644
--- a/tensorflow/compiler/xla/xla_data.proto
+++ b/tensorflow/compiler/xla/xla_data.proto
@@ -100,6 +100,8 @@
// A format specifies the method used by a layout to store an array in memory.
enum Format {
+ // TODO(b/120869032): Rename this to FORMAT_NONE or something else which
+ // better corresponds to its meaning.
INVALID_FORMAT = 0;
// The default layout, with exactly one storage location per element.
DENSE = 1;
@@ -109,8 +111,9 @@
}
// Describes a tile used in tiling-based layout. Refer to
-// g3doc/layout_with_tiling.md for details about tiling-based layout.
-message Tile {
+// g3doc/third_party/tensorflow/compiler/xla/g3doc/layout_with_tiling.md for
+// details about tiling-based layout.
+message TileProto {
// Number of elements in each dimension of the tile. It's ordered from the
// most major dimension of the tile to the most minor dimension of the tile.
// The dimensions correspond to a suffix of the dimensions of the shape being
@@ -128,7 +131,7 @@
// See the XLA documentation for more information on shapes and layouts.
//
// LINT.IfChange
-message Layout {
+message LayoutProto {
// The method used to store the data in memory. The format determines which of
// the other fields are used by the layout.
Format format = 4;
@@ -153,7 +156,7 @@
//
// TODO(b/119839262): implement tiling in each backend or add Unimplemented
// error.
- repeated Tile tiles = 6;
+ repeated TileProto tiles = 6;
// Bit size of each element. If the size is bigger than what the element
// type requires, the value is stored in the least significant
@@ -196,7 +199,7 @@
repeated ShapeProto tuple_shapes = 4;
// The layout used to back this shape.
- Layout layout = 5;
+ LayoutProto layout = 5;
// Important: if any field is added, be sure to modify ShapeUtil::Equal(),
// ShapeUtil::Compatible() and ShapeUtil::Hash() appropriately to account for
diff --git a/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc b/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc
index 3258286..1a5bfac 100644
--- a/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc
+++ b/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc
@@ -120,4 +120,9 @@
.HostMemory("handle"),
XRTReleaseAllocationOp<XRTGenericDeviceAccessor>);
+REGISTER_KERNEL_BUILDER(Name("XRTReleaseAllAllocations").Device(DEVICE_XLA_GPU),
+ XRTReleaseAllAllocationsOp<XRTGenericDeviceAccessor>);
+REGISTER_KERNEL_BUILDER(Name("XRTReleaseAllAllocations").Device(DEVICE_XLA_CPU),
+ XRTReleaseAllAllocationsOp<XRTGenericDeviceAccessor>);
+
} // namespace tensorflow
diff --git a/tensorflow/compiler/xrt/kernels/xrt_state_ops.h b/tensorflow/compiler/xrt/kernels/xrt_state_ops.h
index 26a58fa..e3b292e 100644
--- a/tensorflow/compiler/xrt/kernels/xrt_state_ops.h
+++ b/tensorflow/compiler/xrt/kernels/xrt_state_ops.h
@@ -469,6 +469,26 @@
}
};
+// Op that discards a handle to device memory.
+template <class DeviceAccessor>
+class XRTReleaseAllAllocationsOp : public OpKernel {
+ public:
+ explicit XRTReleaseAllAllocationsOp(OpKernelConstruction* ctx)
+ : OpKernel(ctx) {}
+ ~XRTReleaseAllAllocationsOp() override = default;
+ XRTReleaseAllAllocationsOp(const XRTReleaseAllAllocationsOp&) = delete;
+ XRTReleaseAllAllocationsOp& operator=(const XRTReleaseAllAllocationsOp&) =
+ delete;
+
+ void Compute(OpKernelContext* ctx) override {
+ VLOG(1) << "XRTReleaseAllAllocationsOp::Compute";
+
+ ResourceMgr* rm;
+ OP_REQUIRES_OK(ctx, DeviceAccessor::GetResourceManager(ctx, &rm));
+ OP_REQUIRES_OK(ctx, XRTTupleAllocation::ReleaseAllAllocations(rm));
+ }
+};
+
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_XRT_KERNELS_XRT_STATE_OPS_H_
diff --git a/tensorflow/compiler/xrt/ops/xrt_state_ops.cc b/tensorflow/compiler/xrt/ops/xrt_state_ops.cc
index a3d6310..fe6bee0 100644
--- a/tensorflow/compiler/xrt/ops/xrt_state_ops.cc
+++ b/tensorflow/compiler/xrt/ops/xrt_state_ops.cc
@@ -133,4 +133,11 @@
'handle' is the id returned from the Op that produced the on-device allocation.
)");
+REGISTER_OP("XRTReleaseAllAllocations")
+ .SetShapeFn(tensorflow::shape_inference::NoOutputs)
+ .Doc(
+ R"(
+Discards all the XRT allocations. All the client held handles will be invalid.
+)");
+
} // namespace tensorflow
diff --git a/tensorflow/compiler/xrt/tests/raw_api_test.cc b/tensorflow/compiler/xrt/tests/raw_api_test.cc
index abaa17e..730a227 100644
--- a/tensorflow/compiler/xrt/tests/raw_api_test.cc
+++ b/tensorflow/compiler/xrt/tests/raw_api_test.cc
@@ -265,6 +265,37 @@
&outputs));
}
+TEST(RawApiTest, AllocAndClearAll) {
+ xrt::XLAAllocation alloc;
+ alloc.set_device_ordinal(0);
+ *alloc.mutable_value() =
+ xla::LiteralUtil::CreateR2({{4, 5}, {6, 7}}).ToProto();
+
+ Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag());
+ auto value =
+ ops::Const(root.WithDevice("/device:CPU:0"), alloc.SerializeAsString());
+ auto handle = ops::XRTAllocate(root, value);
+ TF_ASSERT_OK(root.status());
+
+ tensorflow::ClientSession session(root);
+ std::vector<tensorflow::Tensor> outputs;
+ TF_EXPECT_OK(session.Run({handle}, &outputs));
+ EXPECT_EQ(outputs.size(), 1);
+
+ int64 allocation_handle = outputs[0].scalar<int64>()();
+
+ auto clear_all = ops::XRTReleaseAllAllocations(root);
+
+ outputs.clear();
+ TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {},
+ {clear_all}, &outputs));
+ EXPECT_EQ(outputs.size(), 0);
+
+ auto read_after_clear = ops::XRTReadLiteral(root, Input(allocation_handle));
+ EXPECT_EQ(session.Run({read_after_clear}, &outputs).code(),
+ tensorflow::error::Code::NOT_FOUND);
+}
+
TEST(RawApiTest, ReadAndWriteState) {
xrt::XLAAllocation alloc;
alloc.set_device_ordinal(0);
diff --git a/tensorflow/compiler/xrt/xrt_state.cc b/tensorflow/compiler/xrt/xrt_state.cc
index 5c7c537..343460f 100644
--- a/tensorflow/compiler/xrt/xrt_state.cc
+++ b/tensorflow/compiler/xrt/xrt_state.cc
@@ -19,6 +19,7 @@
#include "tensorflow/compiler/xrt/xrt_state.h"
#include <stdint.h>
+#include <map>
#include <memory>
#include <string>
#include <utility>
@@ -34,6 +35,7 @@
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/stream_executor/stream_executor.h"
@@ -41,6 +43,34 @@
namespace {
+class BufferAllocStats {
+ public:
+ struct Stats {
+ int64 count = 0;
+ int64 size = 0;
+ };
+
+ Stats ReportAlloc(int64 device, int64 msize) {
+ mutex_lock lock(lock_);
+ Stats* device_stats = &stats_[device];
+ device_stats->count += 1;
+ device_stats->size += msize;
+ return *device_stats;
+ }
+
+ Stats ReportFree(int64 device, int64 msize) {
+ mutex_lock lock(lock_);
+ Stats* device_stats = &stats_[device];
+ device_stats->count -= 1;
+ device_stats->size -= msize;
+ return *device_stats;
+ }
+
+ private:
+ mutable mutex lock_;
+ std::map<int64, Stats> stats_;
+};
+
const char* kTupleContainer = "tuples";
int64 get_uid() {
@@ -48,6 +78,11 @@
return static_cast<int64>(unsigned_rand);
}
+BufferAllocStats* GetAllocStats() {
+ static BufferAllocStats* stats = new BufferAllocStats();
+ return stats;
+}
+
Status AllocateScopedShapedBuffer(
xla::Backend* backend, int device_ordinal, const xla::Shape& shape,
std::unique_ptr<xla::ScopedShapedBuffer>* buffer) {
@@ -100,9 +135,19 @@
xla::DeviceMemoryAllocator* allocator)
: allocation_(allocation),
device_ordinal_(device_ordinal),
- allocator_(allocator) {}
+ allocator_(allocator) {
+ if (VLOG_IS_ON(2)) {
+ auto stats =
+ GetAllocStats()->ReportAlloc(device_ordinal_, allocation_.size());
+ LOG(INFO) << "XRT Allocation Stats: device=" << device_ordinal_
+ << " count=" << stats.count << " size=" << stats.size;
+ }
+}
XRTBufferAllocation::~XRTBufferAllocation() {
+ if (VLOG_IS_ON(2)) {
+ GetAllocStats()->ReportFree(device_ordinal_, allocation_.size());
+ }
// Deallocate explicitly allows allocation_ to be null.
Status s = allocator_->Deallocate(device_ordinal_, allocation_);
// Nothing to do but check fail here if memory datastructures are corrupted.
@@ -227,6 +272,11 @@
return rm->Delete<XRTTupleAllocation>(kTupleContainer, key_string);
}
+/* static */ Status XRTTupleAllocation::ReleaseAllAllocations(ResourceMgr* rm) {
+ VLOG(1) << "Releasing all XRT held device memory";
+ return rm->Cleanup(kTupleContainer);
+}
+
// Helper typedef to make ShapeTree ForEach helper lambda signatures more
// readable. They need a type of const T& where in this case T is the
// following pointer.
diff --git a/tensorflow/compiler/xrt/xrt_state.h b/tensorflow/compiler/xrt/xrt_state.h
index 3664c0c..3e3d502 100644
--- a/tensorflow/compiler/xrt/xrt_state.h
+++ b/tensorflow/compiler/xrt/xrt_state.h
@@ -129,6 +129,10 @@
// Deletes the reference in the rm to an allocation interned under key.
static Status DeleteFromResourceManager(ResourceMgr* rm, int64 key);
+ // Releases all the device memory allocated by XRT within the resource
+ // manager.
+ static Status ReleaseAllAllocations(ResourceMgr* rm);
+
// Adds the allocation to a ResourceMgr and returns the key that will be used
// to retrieve it. Transfers a reference on *this to rm.
Status Intern(ResourceMgr* rm, int64* key);
diff --git a/tensorflow/contrib/android/cmake/build.gradle b/tensorflow/contrib/android/cmake/build.gradle
index 17a57b9..ddec088 100644
--- a/tensorflow/contrib/android/cmake/build.gradle
+++ b/tensorflow/contrib/android/cmake/build.gradle
@@ -22,8 +22,8 @@
}
externalNativeBuild {
cmake {
- arguments '-DANDROID_TOOLCHAIN=gcc',
- '-DANDROID_STL=gnustl_static'
+ arguments '-DANDROID_TOOLCHAIN=clang',
+ '-DANDROID_STL=c++_static'
}
}
}
@@ -70,7 +70,7 @@
ndkDir = System.getenv('ANDROID_NDK_HOME')
}
-if(! Os.isFamily(Os.FAMILY_WINDOWS)) {
+if (!Os.isFamily(Os.FAMILY_WINDOWS)) {
// This script is for non-Windows OS. For Windows OS, MANUALLY build
// (or copy the built) libs/headers to the
// ${TENSORFLOW_ROOT_DIR}/tensorflow/contrib/makefile/gen
diff --git a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc
index e95dc57..3fe71a2 100644
--- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc
+++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc
@@ -399,6 +399,17 @@
return nullptr;
}
+std::unique_ptr<grpc::ClientAsyncResponseReaderInterface<
+ google::bigtable::v2::CheckAndMutateRowResponse>>
+BigtableTestClient::AsyncCheckAndMutateRow(
+ grpc::ClientContext* context,
+ const google::bigtable::v2::CheckAndMutateRowRequest& request,
+ grpc::CompletionQueue* cq) {
+ LOG(WARNING) << "Call to InMemoryDataClient::" << __func__
+ << "(); this will likely cause a crash!";
+ return nullptr;
+}
+
std::shared_ptr<grpc::Channel> BigtableTestClient::Channel() {
LOG(WARNING) << "Call to InMemoryDataClient::Channel(); this will likely "
"cause a crash!";
diff --git a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h
index c4a1f06..8570590 100644
--- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h
+++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h
@@ -80,6 +80,13 @@
const ::google::bigtable::v2::MutateRowsRequest& request,
::grpc::CompletionQueue* cq, void* tag) override;
+ std::unique_ptr<grpc::ClientAsyncResponseReaderInterface<
+ google::bigtable::v2::CheckAndMutateRowResponse>>
+ AsyncCheckAndMutateRow(
+ grpc::ClientContext* context,
+ const google::bigtable::v2::CheckAndMutateRowRequest& request,
+ grpc::CompletionQueue* cq) override;
+
std::shared_ptr<grpc::Channel> Channel() override;
private:
diff --git a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
index 9f97934..b6cdc7a 100644
--- a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
+++ b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py
@@ -35,8 +35,8 @@
from tensorflow.python.data.experimental.ops import interleave_ops
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.platform import resource_loader
@@ -111,8 +111,7 @@
class BigtableTable(object):
- """BigtableTable is the entrypoint for reading and writing data in Cloud
- Bigtable.
+ """Entry point for reading and writing data in Cloud Bigtable.
This BigtableTable class is the Python representation of the Cloud Bigtable
table within TensorFlow. Methods on this class allow data to be read from and
@@ -593,16 +592,8 @@
self._table = table
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.TensorShape([])
-
- @property
- def output_types(self):
- return dtypes.string
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
class _BigtablePrefixKeyDataset(_BigtableKeyDataset):
@@ -662,16 +653,9 @@
self._columns = [i[1] for i in normalized]
@property
- def output_classes(self):
- return tuple([ops.Tensor] * self._num_outputs)
-
- @property
- def output_shapes(self):
- return tuple([tensor_shape.TensorShape([])] * self._num_outputs)
-
- @property
- def output_types(self):
- return tuple([dtypes.string] * self._num_outputs)
+ def _element_structure(self):
+ return structure.NestedStructure(tuple(
+ [structure.TensorStructure(dtypes.string, [])] * self._num_outputs))
def _as_variant_tensor(self):
# pylint: disable=protected-access
@@ -697,16 +681,9 @@
self._num_outputs = len(normalized) + 1 # 1 for row key
@property
- def output_classes(self):
- return tuple([ops.Tensor] * self._num_outputs)
-
- @property
- def output_shapes(self):
- return tuple([tensor_shape.TensorShape([])] * self._num_outputs)
-
- @property
- def output_types(self):
- return tuple([dtypes.string] * self._num_outputs)
+ def _element_structure(self):
+ return structure.NestedStructure(tuple(
+ [structure.TensorStructure(dtypes.string, [])] * self._num_outputs))
def _as_variant_tensor(self):
return gen_bigtable_ops.bigtable_scan_dataset(
@@ -730,16 +707,10 @@
self._end = end
@property
- def output_classes(self):
- return (ops.Tensor, ops.Tensor)
-
- @property
- def output_shapes(self):
- return (tensor_shape.TensorShape([]), tensor_shape.TensorShape([]))
-
- @property
- def output_types(self):
- return (dtypes.string, dtypes.string)
+ def _element_structure(self):
+ return structure.NestedStructure(
+ (structure.TensorStructure(dtypes.string, []),
+ structure.TensorStructure(dtypes.string, [])))
def _as_variant_tensor(self):
# pylint: disable=protected-access
diff --git a/tensorflow/contrib/data/python/ops/readers.py b/tensorflow/contrib/data/python/ops/readers.py
index aa42782..c015215 100644
--- a/tensorflow/contrib/data/python/ops/readers.py
+++ b/tensorflow/contrib/data/python/ops/readers.py
@@ -21,10 +21,9 @@
from tensorflow.python.data.experimental.ops import readers
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import readers as core_readers
-from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.util import deprecation
@@ -396,18 +395,10 @@
def _as_variant_tensor(self):
return gen_experimental_dataset_ops.experimental_lmdb_dataset(
- self._filenames,
- output_types=nest.flatten(self.output_types),
- output_shapes=nest.flatten(self.output_shapes))
+ self._filenames, **dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return ops.Tensor, ops.Tensor
-
- @property
- def output_shapes(self):
- return (tensor_shape.TensorShape([]), tensor_shape.TensorShape([]))
-
- @property
- def output_types(self):
- return dtypes.string, dtypes.string
+ def _element_structure(self):
+ return structure.NestedStructure(
+ (structure.TensorStructure(dtypes.string, []),
+ structure.TensorStructure(dtypes.string, [])))
diff --git a/tensorflow/contrib/data/python/ops/sliding.py b/tensorflow/contrib/data/python/ops/sliding.py
index 9ebdca3..5c6ee6b 100644
--- a/tensorflow/contrib/data/python/ops/sliding.py
+++ b/tensorflow/contrib/data/python/ops/sliding.py
@@ -39,11 +39,10 @@
self._window_shift = ops.convert_to_tensor(
window_shift, dtype=dtypes.int64, name="window_shift")
- # pylint: disable=protected-access
- input_structure = structure.Structure._from_legacy_structure(
+ input_structure = structure.convert_legacy_structure(
input_dataset.output_types, input_dataset.output_shapes,
input_dataset.output_classes)
- self._output_structure = input_structure._batch(None)
+ self._structure = input_structure._batch(None) # pylint: disable=protected-access
def _as_variant_tensor(self):
return ged_ops.experimental_sliding_window_dataset(
@@ -51,19 +50,11 @@
window_size=self._window_size,
window_shift=self._window_shift,
window_stride=self._window_stride,
- **dataset_ops.flat_structure(structure=self._output_structure))
+ **dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return self._output_structure._to_legacy_output_classes() # pylint: disable=protected-access
-
- @property
- def output_shapes(self):
- return self._output_structure._to_legacy_output_shapes() # pylint: disable=protected-access
-
- @property
- def output_types(self):
- return self._output_structure._to_legacy_output_types() # pylint: disable=protected-access
+ def _element_structure(self):
+ return self._structure
@deprecation.deprecated_args(
diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD
index 4c9c35d..322c02c 100644
--- a/tensorflow/contrib/distribute/python/BUILD
+++ b/tensorflow/contrib/distribute/python/BUILD
@@ -611,6 +611,8 @@
"no_oss", # TODO(b/117919883): Fix python error.
"no_pip",
"no_windows_gpu",
+ # TODO(b/120943676): Re-enable after fixing InvalidArgumentError.
+ "noguitar",
"notsan",
],
)
diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py
index e988b63..346513d 100644
--- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py
+++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py
@@ -70,6 +70,8 @@
self._cross_device_ops = None
self._num_gpus_per_worker = num_gpus_per_worker
self._initialize_local_worker(num_gpus_per_worker)
+ assert isinstance(self._get_cross_device_ops(),
+ cross_device_ops_lib.CollectiveAllReduce)
def _initialize_local_worker(self, num_gpus_per_worker):
"""Initializes the object for local training."""
@@ -77,16 +79,16 @@
self._num_workers = 1
if num_gpus_per_worker:
- local_devices = [
+ local_devices = tuple(
"/device:GPU:%d" % i for i in range(num_gpus_per_worker)
- ]
+ )
else:
- local_devices = ["/device:CPU:0"]
+ local_devices = ("/device:CPU:0",)
self._worker_device = device_util.canonicalize("/device:CPU:0")
self._collective_keys = cross_device_utils.CollectiveKeys()
self._initialize_local(local_devices)
- self._cross_tower_ops = cross_device_ops_lib.CollectiveAllReduce(
+ self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
num_workers=self._num_workers,
num_gpus_per_worker=num_gpus_per_worker,
collective_keys=self._collective_keys)
@@ -104,7 +106,7 @@
if task_type is None or task_id is None:
raise ValueError("When `cluster_spec` is given, you must also specify "
"`task_type` and `task_id`")
- if task_type not in ["chief", "worker"]:
+ if task_type not in ("chief", "worker"):
raise ValueError(
"Unrecognized task_type: %r, valid task types are: \"chief\", "
"\"worker\"." % task_type)
@@ -119,16 +121,16 @@
self._worker_device = "/job:%s/task:%d" % (task_type, task_id)
if num_gpus_per_worker:
- local_devices = [
+ local_devices = tuple(
"%s/device:GPU:%d" % (self._worker_device, i)
for i in range(num_gpus_per_worker)
- ]
+ )
else:
- local_devices = [self._worker_device]
+ local_devices = (self._worker_device,)
self._collective_keys = cross_device_utils.CollectiveKeys()
self._initialize_local(local_devices)
- self._cross_tower_ops = cross_device_ops_lib.CollectiveAllReduce(
+ self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
num_workers=self._num_workers,
num_gpus_per_worker=num_gpus_per_worker,
collective_keys=self._collective_keys)
@@ -267,6 +269,8 @@
# already been initialized with a `cluster_spec`.
self._initialize_multi_worker(self._num_gpus_per_worker, cluster_spec,
task_type, task_id)
+ assert isinstance(self._get_cross_device_ops(),
+ cross_device_ops_lib.CollectiveAllReduce)
if session_config:
session_config.CopyFrom(self._update_config_proto(session_config))
diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py
index 8a9e583..6d7cd14 100644
--- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py
+++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py
@@ -82,7 +82,7 @@
instance_key_with_id_start=num_gpus * 10000 +
CollectiveAllReduceStrategyTestBase.collective_key_base)
distribution.extended._collective_keys = collective_keys
- distribution.extended._inferred_cross_device_ops._collective_keys = (
+ distribution.extended._cross_device_ops._collective_keys = (
collective_keys)
if task_type and task_id is not None:
return distribution, 'grpc://' + self._cluster_spec[task_type][
diff --git a/tensorflow/contrib/distribute/python/estimator_training_test.py b/tensorflow/contrib/distribute/python/estimator_training_test.py
index b369a7f..3f55a8a 100644
--- a/tensorflow/contrib/distribute/python/estimator_training_test.py
+++ b/tensorflow/contrib/distribute/python/estimator_training_test.py
@@ -375,11 +375,13 @@
threads = self.run_multiple_tasks_in_threads(self._independent_worker_fn,
cluster_spec, train_distribute,
eval_distribute)
+ threads_to_join = []
for task_type, ts in threads.items():
if task_type == PS:
continue
for t in ts:
- t.join()
+ threads_to_join.append(t)
+ self.join_independent_workers(threads_to_join)
estimator = self._get_estimator(train_distribute, eval_distribute)
self._inspect_train_and_eval_events(estimator)
@@ -413,8 +415,7 @@
threads = self.run_multiple_tasks_in_threads(self._independent_worker_fn,
cluster_spec, train_distribute,
eval_distribute)
- threads[WORKER][0].join()
- threads[EVALUATOR][0].join()
+ self.join_independent_workers([threads[WORKER][0], threads[EVALUATOR][0]])
estimator = self._get_estimator(train_distribute, eval_distribute)
self._inspect_train_and_eval_events(estimator)
diff --git a/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py b/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py
index 6dfd85b..8c59654 100644
--- a/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py
+++ b/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py
@@ -18,24 +18,12 @@
from __future__ import division
from __future__ import print_function
-import shutil
-import tempfile
from absl.testing import parameterized
import numpy as np
-import six
from tensorflow.contrib.distribute.python import combinations
-from tensorflow.core.protobuf import config_pb2
from tensorflow.python import keras
-from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import distribution_strategy_context as ds_context
-from tensorflow.python.estimator import run_config
-from tensorflow.python.estimator import training
-from tensorflow.python.estimator.canned import dnn_linear_combined
-from tensorflow.python.estimator.canned import prediction_keys
-from tensorflow.python.estimator.export import export
-from tensorflow.python.estimator.inputs import numpy_io
-from tensorflow.python.feature_column import feature_column_lib as feature_column
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
@@ -44,103 +32,7 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
-from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
-from tensorflow.python.summary.writer import writer_cache
-
-
-class KerasOptimizerV2IntegrationTest(test.TestCase, parameterized.TestCase):
-
- def setUp(self):
- self._model_dir = tempfile.mkdtemp()
-
- def dataset_input_fn(self, x, y, batch_size):
-
- def input_fn():
- dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
- dataset = dataset.repeat(1).batch(batch_size)
- return dataset
-
- return input_fn
-
- @combinations.generate(
- combinations.combine(
- mode=['graph'],
- distribution=[
- combinations.one_device_strategy,
- combinations.mirrored_strategy_with_gpu_and_cpu,
- combinations.mirrored_strategy_with_two_gpus,
- combinations.core_mirrored_strategy_with_gpu_and_cpu,
- combinations.core_mirrored_strategy_with_two_gpus
- ],
- use_train_and_evaluate=[True, False]))
- def test_complete_flow_with_mode(self, distribution, use_train_and_evaluate):
- label_dimension = 2
- input_dimension = label_dimension
- batch_size = 10
- data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
- data = data.reshape(batch_size, label_dimension)
- train_input_fn = self.dataset_input_fn(
- x={'x': data},
- y=data,
- batch_size=batch_size // distribution.num_replicas_in_sync)
- eval_input_fn = self.dataset_input_fn(
- x={'x': data},
- y=data,
- batch_size=batch_size // distribution.num_replicas_in_sync)
- predict_input_fn = numpy_io.numpy_input_fn(
- x={'x': data}, batch_size=batch_size, shuffle=False)
-
- linear_feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))
- ]
- dnn_feature_columns = [
- feature_column.numeric_column('x', shape=(input_dimension,))
- ]
- feature_columns = linear_feature_columns + dnn_feature_columns
- session_config = config_pb2.ConfigProto(
- log_device_placement=True, allow_soft_placement=True)
- estimator = dnn_linear_combined.DNNLinearCombinedRegressor(
- linear_feature_columns=linear_feature_columns,
- dnn_hidden_units=(2, 2),
- dnn_feature_columns=dnn_feature_columns,
- label_dimension=label_dimension,
- model_dir=self._model_dir,
- dnn_optimizer=adam.Adam(0.001),
- linear_optimizer=adam.Adam(0.001),
- config=run_config.RunConfig(
- train_distribute=distribution,
- eval_distribute=distribution,
- session_config=session_config))
-
- num_steps = 2
- if use_train_and_evaluate:
- scores, _ = training.train_and_evaluate(
- estimator, training.TrainSpec(train_input_fn, max_steps=num_steps),
- training.EvalSpec(eval_input_fn))
- else:
- estimator.train(train_input_fn, steps=num_steps)
- scores = estimator.evaluate(eval_input_fn)
-
- self.assertIn('loss', six.iterkeys(scores))
-
- predictions = np.array([
- x[prediction_keys.PredictionKeys.PREDICTIONS]
- for x in estimator.predict(predict_input_fn)
- ])
- self.assertAllEqual((batch_size, label_dimension), predictions.shape)
-
- feature_spec = feature_column.make_parse_example_spec(feature_columns)
- serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
- feature_spec)
- export_dir = estimator.export_savedmodel(tempfile.mkdtemp(),
- serving_input_receiver_fn)
- self.assertTrue(gfile.Exists(export_dir))
-
- def tearDown(self):
- if self._model_dir:
- writer_cache.FileWriterCache.clear()
- shutil.rmtree(self._model_dir)
def get_model():
@@ -162,7 +54,9 @@
var = variables.Variable(
2.0, name='var', aggregation=variable_scope.VariableAggregation.SUM)
# grad for cpu is 1, grad for gpu is 2, avg grad is 1.5.
- loss = math_ops.cast(_replica_id() + 1, dtype=dtypes.float32) * var
+ def loss():
+ return math_ops.cast(_replica_id() + 1, dtype=dtypes.float32) * var
+
optimizer = adam.Adam(learning_rate=0.01, beta_1=0.2, beta_2=0.2)
train_op = optimizer.minimize(loss, var_list=[var])
m = optimizer.get_slot(var, 'm')
diff --git a/tensorflow/contrib/distribute/python/keras_test.py b/tensorflow/contrib/distribute/python/keras_test.py
index e530ab6..683cc89 100644
--- a/tensorflow/contrib/distribute/python/keras_test.py
+++ b/tensorflow/contrib/distribute/python/keras_test.py
@@ -165,7 +165,9 @@
return (train_data, test_data)
-def batch_wrapper(dataset, batch_size, distribution):
+def batch_wrapper(dataset, batch_size, distribution, repeat=None):
+ if repeat:
+ dataset = dataset.repeat(repeat)
# TPUs currently require fully defined input shapes, drop_remainder ensures
# the input will have fully defined shapes.
if isinstance(distribution, tpu_strategy.TPUStrategy):
@@ -216,6 +218,7 @@
with_distribution,
x_train, y_train, x_predict):
"""Generates the inputs for correctness check when enable Keras with DS."""
+ training_epochs = 2
global_batch_size = 64
batch_size = global_batch_size
# TODO(b/118776054): Use global batch size for Keras/DS support.
@@ -231,7 +234,7 @@
'batch_size': batch_size,
'x': x_train,
'y': y_train,
- 'epochs': 1,
+ 'epochs': training_epochs,
'shuffle': False,
}
@@ -252,13 +255,14 @@
# keras.fit/evaluate/predict. The batch size is part of the dataset.
train_dataset = dataset_ops.Dataset.from_tensor_slices(
(x_train, y_train))
- x = batch_wrapper(train_dataset, batch_size, with_distribution)
+ x = batch_wrapper(
+ train_dataset, batch_size, with_distribution, repeat=training_epochs)
training_inputs = {
'batch_size': None,
'x': x,
'y': None,
- 'epochs': 1,
+ 'epochs': training_epochs,
'shuffle': False,
'steps_per_epoch': len(x_train) // global_batch_size,
}
@@ -1081,8 +1085,8 @@
@combinations.generate(combinations.combine(
distribution=[
- combinations.mirrored_strategy_with_two_gpus,
- combinations.core_mirrored_strategy_with_two_gpus],
+ combinations.mirrored_strategy_with_gpu_and_cpu,
+ combinations.core_mirrored_strategy_with_gpu_and_cpu],
mode=['graph', 'eager']))
def test_unsupported_features(self, distribution):
with self.cached_session():
@@ -1130,8 +1134,8 @@
@combinations.generate(combinations.combine(
distribution=[
- combinations.mirrored_strategy_with_two_gpus,
- combinations.core_mirrored_strategy_with_two_gpus],
+ combinations.mirrored_strategy_with_gpu_and_cpu,
+ combinations.core_mirrored_strategy_with_gpu_and_cpu],
mode=['graph', 'eager']))
def test_calling_with_unsupported_predefined_callbacks(self, distribution):
with self.cached_session():
@@ -1157,12 +1161,6 @@
'using'):
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0,
callbacks=[keras.callbacks.ReduceLROnPlateau()])
- with self.assertRaisesRegexp(ValueError,
- 'histogram_freq in the TensorBoard callback '
- 'is not supported when using '
- 'DistributionStrategy.'):
- model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0,
- callbacks=[keras.callbacks.TensorBoard(histogram_freq=10)])
class TestDistributionStrategyWithLossMasking(test.TestCase,
@@ -1259,25 +1257,57 @@
train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = batch_wrapper(train_dataset, batch_size, distribution)
- history = model.fit(x=train_dataset, epochs=1, steps_per_epoch=10)
- self.assertEqual(history.history['binary_accuracy'], [1.0])
+ history = model.fit(x=train_dataset, epochs=2, steps_per_epoch=10)
+ self.assertEqual(history.history['binary_accuracy'], [1.0, 1.0])
+
+ @combinations.generate(all_strategy_combinations())
+ def test_eval_metrics_correctness(self, distribution):
+ with self.cached_session():
+ model = keras.Sequential()
+ model.add(
+ keras.layers.Dense(
+ 3, activation='relu', input_dim=4, kernel_initializer='ones'))
+ model.add(
+ keras.layers.Dense(
+ 1, activation='sigmoid', kernel_initializer='ones'))
+ model.compile(
+ loss='mae',
+ metrics=['accuracy', keras.metrics.BinaryAccuracy()],
+ optimizer=gradient_descent.GradientDescentOptimizer(0.001),
+ distribute=distribution)
+
+ # verify correctness of stateful and stateless metrics.
+ x = np.ones((100, 4)).astype('float32')
+ y = np.ones((100, 1)).astype('float32')
+ dataset = dataset_ops.Dataset.from_tensor_slices((x, y)).repeat()
+ dataset = batch_wrapper(dataset, 4, distribution)
+ outs = model.evaluate(dataset, steps=10)
+ self.assertEqual(outs[1], 1.)
+ self.assertEqual(outs[2], 1.)
+
+ y = np.zeros((100, 1)).astype('float32')
+ dataset = dataset_ops.Dataset.from_tensor_slices((x, y)).repeat()
+ dataset = batch_wrapper(dataset, 4, distribution)
+ outs = model.evaluate(dataset, steps=10)
+ self.assertEqual(outs[1], 0.)
+ self.assertEqual(outs[2], 0.)
@combinations.generate(strategy_and_input_combinations())
def test_correctness(self, distribution, use_numpy, use_validation_data):
with self.cached_session():
- tolerance = 1e-5
+ default_tolerance = 1e-5
+ tol_table = {}
if isinstance(distribution, (mirrored_strategy.MirroredStrategy,
mirrored_strategy.CoreMirroredStrategy)):
- # TODO(b/119257215): use the default one once the flakyness is fixed.
- tolerance = 1e-4
-
- if (use_validation_data and
- not isinstance(distribution, tpu_strategy.TPUStrategy)):
- # TODO(b/120435565): Enable tests with use_validation_data once the
- # the underlying bug is fixed.
- return
+ # TODO(b/119257215): Weights are not exactly the same, so use larger
+ # tolerance for now. Predict should be related to weights.
+ tol_table = {
+ 'weights_1': 1e-4,
+ 'weights_2': 1e-4,
+ 'predict_result_1': 1e-4,
+ }
keras.backend.set_image_data_format('channels_last')
np.random.seed(_RANDOM_SEED)
@@ -1318,6 +1348,7 @@
model.compile(
loss=keras.losses.mean_squared_error,
optimizer=gradient_descent_keras.SGD(0.5),
+ metrics=['mse'],
distribute=with_distribution)
training_inputs, eval_inputs, predict_inputs = (
@@ -1325,45 +1356,47 @@
with_distribution,
x_train, y_train, x_predict))
- traning_history = model.fit(**training_inputs).history
+ result = {}
+ result['training_history_1'] = model.fit(**training_inputs).history
if eval_inputs is not None:
- eval_result = model.evaluate(**eval_inputs)
- else:
- # Creates a dummy identical eval_result to be compared later.
- eval_result = 1.0
+ result['eval_result_1'] = model.evaluate(**eval_inputs)
- weights = model.get_weights()
- predict_result = model.predict(**predict_inputs)
+ result['weights_1'] = model.get_weights()
+ result['predict_result_1'] = model.predict(**predict_inputs)
- return weights, traning_history, eval_result, predict_result
+ # Train and eval again to mimic user's flow.
- wts_with_ds, history_with_ds, eval_with_ds, predict_with_ds = (
- fit_eval_and_predict(with_distribution=distribution))
+ result['training_history_2'] = model.fit(**training_inputs).history
- (wts_without_ds, history_without_ds, eval_without_ds,
- predict_without_ds) = fit_eval_and_predict(with_distribution=None)
+ if eval_inputs is not None:
+ result['eval_result_2'] = model.evaluate(**eval_inputs)
+
+ result['weights_2'] = model.get_weights()
+
+ return result
+
+ results_with_ds = fit_eval_and_predict(with_distribution=distribution)
+ results_without_ds = fit_eval_and_predict(with_distribution=None)
# Verify that the weights, training history, eval results, predict outputs
# are the same within some limits of tolerance.
- self.assertAllClose(
- wts_with_ds, wts_without_ds, atol=tolerance, rtol=tolerance,
- msg='Fail to assert weights after training.')
+ for key in results_with_ds:
+ if (key.startswith('training_history') and
+ isinstance(distribution, tpu_strategy.TPUStrategy) and
+ distribution.extended.steps_per_run > 1):
+ # TODO(b/119894254): Enable this test for all cases once the
+ # underlying bug is fixed.
+ continue
- self.assertAllClose(
- eval_with_ds, eval_without_ds, atol=tolerance, rtol=tolerance,
- msg='Fail to assert eval results.')
- self.assertAllClose(
- predict_with_ds, predict_without_ds, atol=tolerance, rtol=tolerance,
- msg='Fail to assert predict results.')
+ tolerance = tol_table.get(key, default_tolerance)
- if not (isinstance(distribution, tpu_strategy.TPUStrategy)
- and distribution.extended.steps_per_run > 1):
- # TODO(b/119894254): Enable this test for all cases once the underlying
- # bug is fixed.
self.assertAllClose(
- history_with_ds, history_without_ds, atol=tolerance, rtol=tolerance,
- msg='Fail to assert training history.')
+ results_with_ds[key],
+ results_without_ds[key],
+ atol=tolerance,
+ rtol=tolerance,
+ msg='Fail to assert {}.'.format(key))
if __name__ == '__main__':
diff --git a/tensorflow/contrib/distribute/python/minimize_loss_test.py b/tensorflow/contrib/distribute/python/minimize_loss_test.py
index dcc9df4..f09483c 100644
--- a/tensorflow/contrib/distribute/python/minimize_loss_test.py
+++ b/tensorflow/contrib/distribute/python/minimize_loss_test.py
@@ -232,7 +232,7 @@
fetches = distribution.unwrap(
distribution.call_for_each_replica(model_fn, args=inputs))
if update_ops_in_cross_replica_mode:
- fetches += ops.get_collection(ops.GraphKeys.UPDATE_OPS)
+ fetches += tuple(ops.get_collection(ops.GraphKeys.UPDATE_OPS))
return control_flow_ops.group(fetches)
iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
@@ -443,7 +443,7 @@
step_fn, iterator, iterations=2,
initial_loop_values=initial_loop_values)
- self.assertEqual({key1: [value1]}, ctx.non_tensor_outputs)
+ self.assertEqual({key1: (value1,)}, ctx.non_tensor_outputs)
self._verify_loss_output(
initial_loss(),
loss_output=ctx.last_step_outputs["replica_loss_reduced"],
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py
index 20f1a08..24399db 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py
@@ -28,7 +28,6 @@
# pylint: disable=protected-access,invalid-name
_call_for_each_replica = mirrored_strategy._call_for_each_replica
-_reduce_non_distributed_value = mirrored_strategy._reduce_non_distributed_value
_create_mirrored_variable = mirrored_strategy._create_mirrored_variable
all_local_devices = mirrored_strategy.all_local_devices
CoreMirroredStrategy = mirrored_strategy.MirroredStrategy
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
index 66512f9..337a86b 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
@@ -180,9 +180,37 @@
variable_scope.variable_creator_scope(main_thread_creator):
result = distribution.extended.call_for_each_replica(model_fn)
result = distribution.unwrap(result)
- expected = ["main_thread:thread_0", "main_thread:thread_1"]
+ expected = ("main_thread:thread_0", "main_thread:thread_1")
self.assertEqual(expected, result)
+@combinations.generate(combinations.combine(
+ distribution=[
+ combinations.mirrored_strategy_with_gpu_and_cpu,
+ combinations.core_mirrored_strategy_with_gpu_and_cpu],
+ mode=["graph", "eager"]))
+class MirroredStrategyCallForEachReplicaTest(test.TestCase):
+
+ def testExecutingEagerlyOutsideFunction(self, distribution):
+ """Verify we preserve the value of executing_eagerly_outside_functions()."""
+ def model_fn():
+ return ops.executing_eagerly_outside_functions()
+
+ originally = ops.executing_eagerly_outside_functions()
+ with distribution.scope():
+ in_scope = ops.executing_eagerly_outside_functions()
+ in_model_fn = distribution.extended.call_for_each_replica(model_fn)
+ unwrapped = distribution.unwrap(in_model_fn)
+ self.assertEqual(in_scope, unwrapped[0])
+ self.assertEqual(in_scope, originally)
+
+ # Verify this all again, but this time in a FuncGraph.
+ with func_graph.FuncGraph("fg").as_default(), distribution.scope():
+ in_scope = ops.executing_eagerly_outside_functions()
+ in_model_fn = distribution.extended.call_for_each_replica(model_fn)
+ unwrapped = distribution.unwrap(in_model_fn)
+ self.assertEqual(in_scope, unwrapped[0])
+ self.assertEqual(in_scope, originally)
+
@combinations.generate(combinations.combine(
distribution=[
diff --git a/tensorflow/contrib/distribute/python/moving_averages_test.py b/tensorflow/contrib/distribute/python/moving_averages_test.py
index c492d8b..8f13e91 100644
--- a/tensorflow/contrib/distribute/python/moving_averages_test.py
+++ b/tensorflow/contrib/distribute/python/moving_averages_test.py
@@ -139,6 +139,27 @@
(2.0 * 0.25 + 0.0) / (1.0 * 0.25 + 1.0)],
var.eval())
+ @combinations.generate(all_combinations)
+ def testAssignVariable(self, distribution):
+
+ def replica_fn():
+ var = variables.Variable([10.0, 11.0])
+ # Here we expect to check the case when input value are variable.
+ val = variables.Variable([1., 2.])
+ decay = 0.25
+ assign = moving_averages.assign_moving_average(
+ var, val, decay, zero_debias=False)
+ return var, assign
+
+ with distribution.scope(), self.cached_session() as sess:
+ var, assign = distribution.call_for_each_replica(replica_fn)
+ variables.global_variables_initializer().run()
+ self.assertAllClose([10.0, 11.0], var.eval())
+ sess.run(distribution.unwrap(assign))
+ self.assertAllClose(
+ [10 * 0.25 + 1. * (1 - 0.25), 11 * 0.25 + 2. * (1 - 0.25)],
+ var.eval())
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/contrib/distribute/python/multi_worker_test_base.py b/tensorflow/contrib/distribute/python/multi_worker_test_base.py
index 147c9b8..b05aac4 100644
--- a/tensorflow/contrib/distribute/python/multi_worker_test_base.py
+++ b/tensorflow/contrib/distribute/python/multi_worker_test_base.py
@@ -40,6 +40,7 @@
from tensorflow.python.estimator import run_config
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import coordinator
from tensorflow.python.training import server_lib
ASSIGNED_PORTS = set()
@@ -360,6 +361,7 @@
self._mock_os_env = MockOsEnv()
self._mock_context = test.mock.patch.object(os, 'environ',
self._mock_os_env)
+ self._coord = coordinator.Coordinator()
super(IndependentWorkerTestBase, self).setUp()
self._mock_context.__enter__()
@@ -368,8 +370,9 @@
super(IndependentWorkerTestBase, self).tearDown()
def _task_thread(self, task_fn, tf_config, *args, **kwargs):
- os.environ['TF_CONFIG'] = json.dumps(tf_config)
- task_fn(*args, **kwargs)
+ with self._coord.stop_on_exception():
+ os.environ['TF_CONFIG'] = json.dumps(tf_config)
+ task_fn(*args, **kwargs)
def _run_task_in_thread(self, task_fn, cluster_spec, task_type, task_id,
*args, **kwargs):
@@ -403,3 +406,6 @@
*args, **kwargs)
threads[task_type].append(t)
return threads
+
+ def join_independent_workers(self, worker_threads):
+ self._coord.join(worker_threads)
diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py
index e322b6a..fdbfba4 100644
--- a/tensorflow/contrib/distribute/python/one_device_strategy.py
+++ b/tensorflow/contrib/distribute/python/one_device_strategy.py
@@ -60,7 +60,7 @@
if isinstance(colocate_with, six.string_types):
with ops.device(colocate_with):
return next_creator(*args, **kwargs)
- if (isinstance(colocate_with, list) and len(colocate_with) == 1 and
+ if (isinstance(colocate_with, (list, tuple)) and len(colocate_with) == 1 and
isinstance(colocate_with[0], six.string_types)):
with ops.device(colocate_with[0]):
return next_creator(*args, **kwargs)
@@ -166,7 +166,7 @@
return array_ops.identity(replica_local_var)
def _unwrap(self, value):
- return [value]
+ return (value,)
def value_container(self, value):
return value
@@ -177,15 +177,15 @@
@property
def worker_devices(self):
- return [self._device]
+ return (self._device,)
@property
def parameter_devices(self):
- return [self._device]
+ return (self._device,)
def non_slot_devices(self, var_list):
del var_list
- return [self._device]
+ return (self._device,)
@property
def experimental_should_init(self):
@@ -216,4 +216,4 @@
@property
def devices(self):
- return [self._distribution_strategy.extended.worker_devices[0]]
+ return self._distribution_strategy.extended.worker_devices
diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy.py b/tensorflow/contrib/distribute/python/parameter_server_strategy.py
index eaeb4d7..ca51b07 100644
--- a/tensorflow/contrib/distribute/python/parameter_server_strategy.py
+++ b/tensorflow/contrib/distribute/python/parameter_server_strategy.py
@@ -145,14 +145,14 @@
# replica. When there are GPUs, replicate operations on these GPUs.
# Otherwise, place operations on CPU.
if num_gpus_per_worker > 0:
- self._compute_devices = [
+ self._compute_devices = tuple(
"%s/device:GPU:%d" % (self._worker_device, i)
for i in range(num_gpus_per_worker)
- ]
+ )
else:
- self._compute_devices = [self._worker_device]
+ self._compute_devices = (self._worker_device,)
- self._compute_devices = list(
+ self._compute_devices = tuple(
map(device_util.resolve, self._compute_devices))
self._canonical_compute_device_set = set(self._compute_devices)
@@ -176,8 +176,8 @@
# The `_parameter_devices` is needed for the `parameter_devices` property
# and is a list of all variable devices. Here parameter devices are all
# tasks of the "ps" job.
- self._parameter_devices = map("/job:ps/task:{}".format,
- range(num_ps_replicas))
+ self._parameter_devices = tuple(map("/job:ps/task:{}".format,
+ range(num_ps_replicas)))
# Add a default device so that ops without specified devices will not end up
# on other workers.
@@ -204,24 +204,24 @@
# replica. When there are GPUs, replicate operations on these GPUs.
# Otherwise, place operations on CPU.
if num_gpus_per_worker > 0:
- self._compute_devices = list(
+ self._compute_devices = tuple(
map("/device:GPU:{}".format, range(num_gpus_per_worker)))
else:
- self._compute_devices = [_LOCAL_CPU]
+ self._compute_devices = (_LOCAL_CPU,)
- self._compute_devices = list(
+ self._compute_devices = tuple(
map(device_util.resolve, self._compute_devices))
self._canonical_compute_device_set = set(self._compute_devices)
# If there is only one GPU, put everything on that GPU. Otherwise, place
# variables on CPU.
if num_gpus_per_worker == 1:
- assert len(list(self._compute_devices)) == 1
+ assert len(self._compute_devices) == 1
self._variable_device = _LOCAL_GPU_0
- self._parameter_devices = [_LOCAL_GPU_0]
+ self._parameter_devices = (_LOCAL_GPU_0,)
else:
self._variable_device = _LOCAL_CPU
- self._parameter_devices = [_LOCAL_CPU]
+ self._parameter_devices = (_LOCAL_CPU,)
self._is_chief = True
self._cluster_spec = None
@@ -356,7 +356,7 @@
self._verify_destinations_not_different_worker(destinations)
if not isinstance(value, values.DistributedValues):
# pylint: disable=protected-access
- return mirrored_strategy._reduce_non_distributed_value(
+ return cross_device_ops_lib.reduce_non_distributed_value(
self, reduce_op, value, destinations)
return self._cross_device_ops.reduce(
reduce_op, value, destinations=destinations)
@@ -417,9 +417,9 @@
if isinstance(val, values.DistributedValues):
# Return in a deterministic order.
if set(val.devices) == self._canonical_compute_device_set:
- return [val.get(device=d) for d in self._compute_devices]
- return [val.get(device=d) for d in sorted(val.devices)]
- return [val]
+ return tuple(val.get(device=d) for d in self._compute_devices)
+ return tuple(val.get(device=d) for d in sorted(val.devices))
+ return (val,)
def value_container(self, val):
if (hasattr(val, "_aggregating_container") and
@@ -497,12 +497,11 @@
@property
def worker_devices(self):
- # Make a copy to prevent users from accidentally mutating our copy.
- return list(self._compute_devices)
+ return self._compute_devices
@property
def parameter_devices(self):
- return list(self._parameter_devices)
+ return self._parameter_devices
def non_slot_devices(self, var_list):
return min(var_list, key=lambda x: x.name)
diff --git a/tensorflow/contrib/distribute/python/strategy_test_lib.py b/tensorflow/contrib/distribute/python/strategy_test_lib.py
index d50b142..d441b5a 100644
--- a/tensorflow/contrib/distribute/python/strategy_test_lib.py
+++ b/tensorflow/contrib/distribute/python/strategy_test_lib.py
@@ -290,4 +290,4 @@
self.evaluate(strategy.group(train_ops))
global_step_tensors = strategy.unwrap(value)
global_step_values = self.evaluate(global_step_tensors)
- self.assertEqual([1] * len(global_step_tensors), global_step_values)
+ self.assertEqual((1,) * len(global_step_tensors), global_step_values)
diff --git a/tensorflow/contrib/distribute/python/tpu_strategy.py b/tensorflow/contrib/distribute/python/tpu_strategy.py
index 39ed8f7..7ea245e 100644
--- a/tensorflow/contrib/distribute/python/tpu_strategy.py
+++ b/tensorflow/contrib/distribute/python/tpu_strategy.py
@@ -28,6 +28,8 @@
from tensorflow.contrib.tpu.python.tpu import tpu
from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib
from tensorflow.contrib.tpu.python.tpu import training_loop
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session as session_lib
from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import distribute_lib
@@ -43,12 +45,10 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_scope as vs
+from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import nest
-_TPU_INITIALIZE_SYSTEM_COLLECTION = "TPU_STRATEGY_INITIALIZE"
-
-
def get_tpu_system_metadata(tpu_cluster_resolver):
"""Retrieves TPU system metadata given a TPUClusterResolver."""
master = tpu_cluster_resolver.master()
@@ -145,6 +145,9 @@
class TPUExtended(distribute_lib.DistributionStrategyExtended):
"""Implementation of TPUStrategy."""
+ # Track what TPU devices have been initialized.
+ _initialized_devices = []
+
def __init__(self, container_strategy, tpu_cluster_resolver, steps_per_run,
num_cores=None):
super(TPUExtended, self).__init__(container_strategy)
@@ -159,16 +162,41 @@
if "device:TPU:" in d.name}
self._device_index = values.PerReplica(device_map)
self._host_device = self.get_host_cpu_device(0)
- self._tpu_devices = sorted(device_map.keys())
+ self._tpu_devices = tuple(sorted(device_map.keys()))
# Only create variables for the number of replicas we're running.
self._tpu_devices = self._tpu_devices[:self._num_replicas_in_sync]
# TODO(sourabhbajaj): Remove this once performance of running one step
# at a time is comparable to multiple steps.
self.steps_per_run = steps_per_run
-
self._require_static_shapes = True
+ # Initialize the TPU devices.
+ self._initialize_tpu()
+
+ def _initialize_tpu(self):
+ """Initialize the TPU devices in a separate session and graph.
+
+ We keep track of all the TPU devices that we're initialized as we should
+ only be running TPU initialize once for the entire process.
+ """
+ master = self._tpu_cluster_resolver.master()
+ # Verify TPU has not already been initialized in this process.
+ if master in TPUExtended._initialized_devices:
+ logging.info("TPU master %s has already been initialized." % master)
+ return
+
+ logging.info("Initializing the TPU system.")
+ session_config = config_pb2.ConfigProto(allow_soft_placement=True)
+ self._configure(session_config)
+ with ops.Graph().as_default():
+ with session_lib.Session(config=session_config, target=master) as sess:
+ sess.run([tpu.initialize_system()])
+ logging.info("Finized initializing TPU system.")
+
+ # Update Strategy state to make sure we can track device initialization.
+ TPUExtended._initialized_devices.append(master)
+
def _get_enqueue_op_per_host(self, host_id, multi_worker_iterator,
input_shapes, iterations):
"""Create an enqueue op for a single host identified using host_id.
@@ -380,22 +408,14 @@
# TODO(priyag): Add appopriate call here when eager is supported for TPUs.
raise NotImplementedError("Eager mode not supported in TPUStrategy.")
else:
- # TODO(jhseu): We need this hack because DistributionStrategies must be
- # pickleable for copy.deepcopy(). Remove when initialize_system goes away.
- graph = ops.get_default_graph()
- tpu_init = graph.get_collection(_TPU_INITIALIZE_SYSTEM_COLLECTION)
- if tpu_init:
- return tpu_init
- graph.add_to_collection(_TPU_INITIALIZE_SYSTEM_COLLECTION,
- tpu.initialize_system())
- return graph.get_collection(_TPU_INITIALIZE_SYSTEM_COLLECTION)
+ return []
def _finalize(self):
if context.executing_eagerly():
# TODO(priyag): Add appopriate call here when eager is supported for TPUs.
raise NotImplementedError("Eager mode not supported in TPUStrategy.")
else:
- return [tpu.shutdown_system()]
+ return []
def _get_devices_from(self, colocate_with=None):
# TODO(jhseu): Change this when we support model parallelism.
@@ -445,6 +465,14 @@
"Currently only support sum & mean in TPUStrategy.")
return tpu_ops.cross_replica_sum(value)
+ if not isinstance(value, values.DistributedValues):
+ # This function handles reducing values that are not PerReplica or
+ # Mirrored values. For example, the same value could be present on all
+ # replicas in which case `value` would be a single value or value could
+ # be 0.
+ return cross_device_ops_lib.reduce_non_distributed_value(
+ self, reduce_op, value, destinations)
+
# Validate that the destination is same as the host device
# Note we don't do this when in replicate context as the reduction is
# performed on the TPU device itself.
@@ -487,13 +515,13 @@
def _unwrap(self, val):
if isinstance(val, values.DistributedValues):
# Return in a deterministic order.
- return [val.get(device=d) for d in sorted(val.devices)]
+ return tuple(val.get(device=d) for d in sorted(val.devices))
elif isinstance(val, list):
# TODO(josh11b): We need to remove this case; per device values should
# be represented using a PerReplica wrapper instead of a list with
# one entry per device.
- return val
- return [val]
+ return tuple(val)
+ return (val,)
def value_container(self, value):
return value
@@ -599,4 +627,4 @@
distribute_lib.require_replica_context(self)
ds = self._distribution_strategy
replica_id = tensor_util.constant_value(self._replica_id_in_sync_group)
- return [ds.extended.worker_devices[replica_id]]
+ return (ds.extended.worker_devices[replica_id],)
diff --git a/tensorflow/contrib/eager/python/tfe.py b/tensorflow/contrib/eager/python/tfe.py
index 33c988f..8882a86 100644
--- a/tensorflow/contrib/eager/python/tfe.py
+++ b/tensorflow/contrib/eager/python/tfe.py
@@ -41,6 +41,8 @@
@@add_execution_callback
@@clear_execution_callbacks
+@@errstate
+@@ExecutionCallback
@@inf_callback
@@inf_nan_callback
@@nan_callback
@@ -119,6 +121,8 @@
from tensorflow.python.eager.def_function import function
from tensorflow.python.eager.execution_callbacks import add_execution_callback
from tensorflow.python.eager.execution_callbacks import clear_execution_callbacks
+from tensorflow.python.eager.execution_callbacks import errstate
+from tensorflow.python.eager.execution_callbacks import ExecutionCallback
from tensorflow.python.eager.execution_callbacks import inf_callback
from tensorflow.python.eager.execution_callbacks import inf_nan_callback
from tensorflow.python.eager.execution_callbacks import nan_callback
diff --git a/tensorflow/contrib/feature_column/BUILD b/tensorflow/contrib/feature_column/BUILD
index 1cd83bd..4c1d1a2 100644
--- a/tensorflow/contrib/feature_column/BUILD
+++ b/tensorflow/contrib/feature_column/BUILD
@@ -110,8 +110,8 @@
"//tensorflow/python:parsing_ops",
"//tensorflow/python:sparse_tensor",
"//tensorflow/python:training",
- "//tensorflow/python/feature_column",
"//tensorflow/python/feature_column:feature_column_py",
+ "//tensorflow/python/feature_column:feature_column_v2_test",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py
index 0d34ad1..83b93ec 100644
--- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py
+++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py
@@ -203,7 +203,8 @@
columns = [watches_embedding]
features = tf.parse_example(..., features=make_parse_example_spec(columns))
- input_layer, sequence_length = sequence_input_layer(features, columns)
+ sequence_feature_layer = SequenceFeatureLayer(columns)
+ input_layer, sequence_length = sequence_feature_layer(features)
rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size)
outputs, state = tf.nn.dynamic_rnn(
@@ -219,15 +220,17 @@
`[0, num_buckets)`, and will replace out-of-range inputs.
Returns:
- A `_SequenceCategoricalColumn`.
+ A `SequenceCategoricalColumn`.
Raises:
ValueError: if `num_buckets` is less than one.
ValueError: if `default_value` is not in range `[0, num_buckets)`.
"""
- return fc_old._SequenceCategoricalColumn(
- fc_old._categorical_column_with_identity(
- key=key, num_buckets=num_buckets, default_value=default_value))
+ return fc.SequenceCategoricalColumn(
+ fc.categorical_column_with_identity(
+ key=key,
+ num_buckets=num_buckets,
+ default_value=default_value))
def sequence_categorical_column_with_hash_bucket(
@@ -247,7 +250,8 @@
columns = [tokens_embedding]
features = tf.parse_example(..., features=make_parse_example_spec(columns))
- input_layer, sequence_length = sequence_input_layer(features, columns)
+ sequence_feature_layer = SequenceFeatureLayer(columns)
+ input_layer, sequence_length = sequence_feature_layer(features)
rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size)
outputs, state = tf.nn.dynamic_rnn(
@@ -260,15 +264,17 @@
dtype: The type of features. Only string and integer types are supported.
Returns:
- A `_SequenceCategoricalColumn`.
+ A `SequenceCategoricalColumn`.
Raises:
ValueError: `hash_bucket_size` is not greater than 1.
ValueError: `dtype` is neither string nor integer.
"""
- return fc_old._SequenceCategoricalColumn(
- fc_old._categorical_column_with_hash_bucket(
- key=key, hash_bucket_size=hash_bucket_size, dtype=dtype))
+ return fc.SequenceCategoricalColumn(
+ fc.categorical_column_with_hash_bucket(
+ key=key,
+ hash_bucket_size=hash_bucket_size,
+ dtype=dtype))
def sequence_categorical_column_with_vocabulary_file(
@@ -290,7 +296,8 @@
columns = [states_embedding]
features = tf.parse_example(..., features=make_parse_example_spec(columns))
- input_layer, sequence_length = sequence_input_layer(features, columns)
+ sequence_feature_layer = SequenceFeatureLayer(columns)
+ input_layer, sequence_length = sequence_feature_layer(features)
rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size)
outputs, state = tf.nn.dynamic_rnn(
@@ -314,7 +321,7 @@
dtype: The type of features. Only string and integer types are supported.
Returns:
- A `_SequenceCategoricalColumn`.
+ A `SequenceCategoricalColumn`.
Raises:
ValueError: `vocabulary_file` is missing or cannot be opened.
@@ -323,8 +330,8 @@
ValueError: `num_oov_buckets` and `default_value` are both specified.
ValueError: `dtype` is neither string nor integer.
"""
- return fc_old._SequenceCategoricalColumn(
- fc_old._categorical_column_with_vocabulary_file(
+ return fc.SequenceCategoricalColumn(
+ fc.categorical_column_with_vocabulary_file(
key=key,
vocabulary_file=vocabulary_file,
vocabulary_size=vocabulary_size,
@@ -351,7 +358,8 @@
columns = [colors_embedding]
features = tf.parse_example(..., features=make_parse_example_spec(columns))
- input_layer, sequence_length = sequence_input_layer(features, columns)
+ sequence_feature_layer = SequenceFeatureLayer(columns)
+ input_layer, sequence_length = sequence_feature_layer(features)
rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size)
outputs, state = tf.nn.dynamic_rnn(
@@ -375,7 +383,7 @@
with `default_value`.
Returns:
- A `_SequenceCategoricalColumn`.
+ A `SequenceCategoricalColumn`.
Raises:
ValueError: if `vocabulary_list` is empty, or contains duplicate keys.
@@ -383,8 +391,8 @@
ValueError: `num_oov_buckets` and `default_value` are both specified.
ValueError: if `dtype` is not integer or string.
"""
- return fc_old._SequenceCategoricalColumn(
- fc_old._categorical_column_with_vocabulary_list(
+ return fc.SequenceCategoricalColumn(
+ fc.categorical_column_with_vocabulary_list(
key=key,
vocabulary_list=vocabulary_list,
dtype=dtype,
diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py
index ca4398a..be012a8 100644
--- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py
+++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py
@@ -26,7 +26,7 @@
from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column_v2 as sfc
from tensorflow.python.feature_column import feature_column as fc_old
from tensorflow.python.feature_column import feature_column_lib as fc
-from tensorflow.python.feature_column.feature_column import _LazyBuilder
+from tensorflow.python.feature_column.feature_column_v2_test import _TestStateManager
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
@@ -131,7 +131,7 @@
feature_columns=[embedding_column_b, embedding_column_a])
global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
- self.assertItemsEqual(
+ self.assertCountEqual(
('sequence_input_layer/aaa_embedding/embedding_weights:0',
'sequence_input_layer/bbb_embedding/embedding_weights:0'),
tuple([v.name for v in global_vars]))
@@ -223,7 +223,7 @@
feature_columns=shared_embedding_columns)
global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
- self.assertItemsEqual(
+ self.assertCountEqual(
('sequence_input_layer/aaa_bbb_shared_embedding/embedding_weights:0',),
tuple([v.name for v in global_vars]))
with monitored_session.MonitoredSession() as sess:
@@ -670,6 +670,23 @@
test_case.assertAllEqual(expected.dense_shape, actual.dense_shape)
+def _get_sequence_dense_tensor(column, features):
+ return column.get_sequence_dense_tensor(
+ fc.FeatureTransformationCache(features), None)
+
+
+def _get_sequence_dense_tensor_state(column, features):
+ state_manager = _TestStateManager()
+ column.create_state(state_manager)
+ return column.get_sequence_dense_tensor(
+ fc.FeatureTransformationCache(features), state_manager)
+
+
+def _get_sparse_tensors(column, features):
+ return column.get_sparse_tensors(
+ fc.FeatureTransformationCache(features), None)
+
+
class SequenceCategoricalColumnWithIdentityTest(
test.TestCase, parameterized.TestCase):
@@ -698,7 +715,7 @@
expected = sparse_tensor.SparseTensorValue(**expected_args)
column = sfc.sequence_categorical_column_with_identity('aaa', num_buckets=9)
- id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs}))
+ id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs})
self.assertIsNone(id_weight_pair.weight_tensor)
with monitored_session.MonitoredSession() as sess:
@@ -737,7 +754,7 @@
column = sfc.sequence_categorical_column_with_hash_bucket(
'aaa', hash_bucket_size=10)
- id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs}))
+ id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs})
self.assertIsNone(id_weight_pair.weight_tensor)
with monitored_session.MonitoredSession() as sess:
@@ -790,7 +807,7 @@
vocabulary_file=self._wire_vocabulary_file_name,
vocabulary_size=self._wire_vocabulary_size)
- id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs}))
+ id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs})
self.assertIsNone(id_weight_pair.weight_tensor)
with monitored_session.MonitoredSession() as sess:
@@ -814,8 +831,7 @@
input_placeholder_shape[1] = None
input_placeholder = array_ops.sparse_placeholder(
dtypes.string, shape=input_placeholder_shape)
- id_weight_pair = column._get_sparse_tensors(
- _LazyBuilder({'aaa': input_placeholder}))
+ id_weight_pair = _get_sparse_tensors(column, {'aaa': input_placeholder})
self.assertIsNone(id_weight_pair.weight_tensor)
with monitored_session.MonitoredSession() as sess:
@@ -855,7 +871,7 @@
key='aaa',
vocabulary_list=('omar', 'stringer', 'marlo'))
- id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs}))
+ id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs})
self.assertIsNone(id_weight_pair.weight_tensor)
with monitored_session.MonitoredSession() as sess:
@@ -922,13 +938,12 @@
categorical_column = sfc.sequence_categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- embedding_column = fc_old._embedding_column(
- categorical_column,
- dimension=embedding_dimension,
+ embedding_column = fc.embedding_column(
+ categorical_column, dimension=embedding_dimension,
initializer=_initializer)
- embedding_lookup, _ = embedding_column._get_sequence_dense_tensor(
- _LazyBuilder({'aaa': inputs}))
+ embedding_lookup, _ = _get_sequence_dense_tensor_state(
+ embedding_column, {'aaa': inputs})
global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
self.assertItemsEqual(
@@ -961,10 +976,11 @@
categorical_column = sfc.sequence_categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- embedding_column = fc_old._embedding_column(categorical_column, dimension=2)
+ embedding_column = fc.embedding_column(
+ categorical_column, dimension=2)
- _, sequence_length = embedding_column._get_sequence_dense_tensor(
- _LazyBuilder({'aaa': inputs}))
+ _, sequence_length = _get_sequence_dense_tensor_state(
+ embedding_column, {'aaa': inputs})
with monitored_session.MonitoredSession() as sess:
sequence_length = sess.run(sequence_length)
@@ -988,10 +1004,11 @@
categorical_column = sfc.sequence_categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- embedding_column = fc_old._embedding_column(categorical_column, dimension=2)
+ embedding_column = fc.embedding_column(
+ categorical_column, dimension=2)
- _, sequence_length = embedding_column._get_sequence_dense_tensor(
- _LazyBuilder({'aaa': sparse_input}))
+ _, sequence_length = _get_sequence_dense_tensor_state(
+ embedding_column, {'aaa': sparse_input})
with monitored_session.MonitoredSession() as sess:
self.assertAllEqual(
@@ -1058,22 +1075,18 @@
key='aaa', num_buckets=vocabulary_size)
categorical_column_b = sfc.sequence_categorical_column_with_identity(
key='bbb', num_buckets=vocabulary_size)
- shared_embedding_columns = fc.shared_embedding_columns(
+ shared_embedding_columns = fc.shared_embedding_columns_v2(
[categorical_column_a, categorical_column_b],
dimension=embedding_dimension,
initializer=_initializer)
- embedding_lookup_a = shared_embedding_columns[0]._get_sequence_dense_tensor(
- _LazyBuilder({
- 'aaa': sparse_input_a
- }))[0]
- embedding_lookup_b = shared_embedding_columns[1]._get_sequence_dense_tensor(
- _LazyBuilder({
- 'bbb': sparse_input_b
- }))[0]
+ embedding_lookup_a = _get_sequence_dense_tensor(
+ shared_embedding_columns[0], {'aaa': sparse_input_a})[0]
+ embedding_lookup_b = _get_sequence_dense_tensor(
+ shared_embedding_columns[1], {'bbb': sparse_input_b})[0]
global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
- self.assertItemsEqual(('embedding_weights:0',),
+ self.assertItemsEqual(('aaa_bbb_shared_embedding:0',),
tuple([v.name for v in global_vars]))
with monitored_session.MonitoredSession() as sess:
self.assertAllEqual(embedding_values, global_vars[0].eval(session=sess))
@@ -1104,17 +1117,13 @@
expected_sequence_length_b = [2, 1]
categorical_column_b = sfc.sequence_categorical_column_with_identity(
key='bbb', num_buckets=vocabulary_size)
- shared_embedding_columns = fc.shared_embedding_columns(
+ shared_embedding_columns = fc.shared_embedding_columns_v2(
[categorical_column_a, categorical_column_b], dimension=2)
- sequence_length_a = shared_embedding_columns[0]._get_sequence_dense_tensor(
- _LazyBuilder({
- 'aaa': sparse_input_a
- }))[1]
- sequence_length_b = shared_embedding_columns[1]._get_sequence_dense_tensor(
- _LazyBuilder({
- 'bbb': sparse_input_b
- }))[1]
+ sequence_length_a = _get_sequence_dense_tensor(
+ shared_embedding_columns[0], {'aaa': sparse_input_a})[1]
+ sequence_length_b = _get_sequence_dense_tensor(
+ shared_embedding_columns[1], {'bbb': sparse_input_b})[1]
with monitored_session.MonitoredSession() as sess:
sequence_length_a = sess.run(sequence_length_a)
@@ -1155,17 +1164,13 @@
categorical_column_b = sfc.sequence_categorical_column_with_identity(
key='bbb', num_buckets=vocabulary_size)
- shared_embedding_columns = fc.shared_embedding_columns(
+ shared_embedding_columns = fc.shared_embedding_columns_v2(
[categorical_column_a, categorical_column_b], dimension=2)
- sequence_length_a = shared_embedding_columns[0]._get_sequence_dense_tensor(
- _LazyBuilder({
- 'aaa': sparse_input_a
- }))[1]
- sequence_length_b = shared_embedding_columns[1]._get_sequence_dense_tensor(
- _LazyBuilder({
- 'bbb': sparse_input_b
- }))[1]
+ sequence_length_a = _get_sequence_dense_tensor(
+ shared_embedding_columns[0], {'aaa': sparse_input_a})[1]
+ sequence_length_b = _get_sequence_dense_tensor(
+ shared_embedding_columns[1], {'bbb': sparse_input_b})[1]
with monitored_session.MonitoredSession() as sess:
self.assertAllEqual(
@@ -1221,10 +1226,10 @@
categorical_column = sfc.sequence_categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- indicator_column = fc_old._indicator_column(categorical_column)
+ indicator_column = fc.indicator_column(categorical_column)
- indicator_tensor, _ = indicator_column._get_sequence_dense_tensor(
- _LazyBuilder({'aaa': inputs}))
+ indicator_tensor, _ = _get_sequence_dense_tensor(
+ indicator_column, {'aaa': inputs})
with monitored_session.MonitoredSession() as sess:
self.assertAllEqual(expected, indicator_tensor.eval(session=sess))
@@ -1253,10 +1258,10 @@
categorical_column = sfc.sequence_categorical_column_with_identity(
key='aaa', num_buckets=vocabulary_size)
- indicator_column = fc_old._indicator_column(categorical_column)
+ indicator_column = fc.indicator_column(categorical_column)
- _, sequence_length = indicator_column._get_sequence_dense_tensor(
- _LazyBuilder({'aaa': inputs}))
+ _, sequence_length = _get_sequence_dense_tensor(
+ indicator_column, {'aaa': inputs})
with monitored_session.MonitoredSession() as sess:
sequence_length = sess.run(sequence_length)
@@ -1282,19 +1287,14 @@
key='aaa', num_buckets=vocabulary_size)
indicator_column = fc.indicator_column(categorical_column)
- _, sequence_length = indicator_column._get_sequence_dense_tensor(
- _LazyBuilder({'aaa': sparse_input}))
+ _, sequence_length = _get_sequence_dense_tensor(
+ indicator_column, {'aaa': sparse_input})
with monitored_session.MonitoredSession() as sess:
self.assertAllEqual(
expected_sequence_length, sequence_length.eval(session=sess))
-def _get_sequence_dense_tensor(column, features):
- return column.get_sequence_dense_tensor(
- fc.FeatureTransformationCache(features), None)
-
-
class SequenceNumericColumnTest(test.TestCase, parameterized.TestCase):
def test_defaults(self):
diff --git a/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc b/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc
index 93b1aaa..c541c71 100644
--- a/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc
+++ b/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc
@@ -522,7 +522,7 @@
auto bias_ptr = AsDeviceMemory(bias.template flat<BiasType>().data(),
bias.template flat<BiasType>().size());
- static int64 ConvolveScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveScratchSize = GetDnnWorkspaceLimit(
// default value is in bytes despite the name of the environment variable
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB
);
@@ -570,7 +570,7 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
dnn::ProfileResult profile_result;
bool cudnn_launch_status =
stream
@@ -609,7 +609,7 @@
algorithm_config);
}
- CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
bool cudnn_launch_status =
stream
->ThenFusedConvolveWithAlgorithm(
diff --git a/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py b/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py
index e2594fa..364fa4e 100644
--- a/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py
+++ b/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py
@@ -64,6 +64,9 @@
"""
tensor.shape[1:].assert_is_fully_defined()
num_features = tensor.shape[1:].num_elements()
+ if conditioning.shape.ndims < 2:
+ raise ValueError('conditioning must be at least 2D, but saw shape: %s'
+ % conditioning.shape)
mapped_conditioning = layers.linear(
layers.flatten(conditioning), num_features)
diff --git a/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py b/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py
index 0aad769..f5c7d53 100644
--- a/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py
+++ b/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py
@@ -45,7 +45,7 @@
array_ops.placeholder(dtypes.float32, (5, None)),
array_ops.placeholder(dtypes.float32, (5, 1)))
- with self.assertRaisesRegexp(ValueError, 'expected min_ndim=2'):
+ with self.assertRaisesRegexp(ValueError, 'at least 2D'):
conditioning_utils.condition_tensor(
array_ops.placeholder(dtypes.float32, (5, 2)),
array_ops.placeholder(dtypes.float32, (5)))
diff --git a/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py b/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py
index d3fcc8c..5c55998 100644
--- a/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py
+++ b/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py
@@ -20,10 +20,9 @@
from tensorflow.contrib.hadoop.python.ops import gen_dataset_ops
from tensorflow.contrib.hadoop.python.ops import hadoop_op_loader # pylint: disable=unused-import
from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
class SequenceFileDataset(dataset_ops.DatasetSource):
@@ -57,16 +56,10 @@
def _as_variant_tensor(self):
return gen_dataset_ops.sequence_file_dataset(
- self._filenames, nest.flatten(self.output_types))
+ self._filenames, self._element_structure._flat_types) # pylint: disable=protected-access
@property
- def output_classes(self):
- return ops.Tensor, ops.Tensor
-
- @property
- def output_shapes(self):
- return (tensor_shape.TensorShape([]), tensor_shape.TensorShape([]))
-
- @property
- def output_types(self):
- return dtypes.string, dtypes.string
+ def _element_structure(self):
+ return structure.NestedStructure(
+ (structure.TensorStructure(dtypes.string, []),
+ structure.TensorStructure(dtypes.string, [])))
diff --git a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
index 936b29a..e4762c9 100644
--- a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
+++ b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py
@@ -27,6 +27,7 @@
from tensorflow.contrib.ignite.python.ops import gen_dataset_ops
from tensorflow.contrib.ignite.python.ops import ignite_op_loader # pylint: disable=unused-import
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
@@ -34,10 +35,7 @@
@six.add_metaclass(abc.ABCMeta)
class Readable(object):
- """Readable abstract class that exposes methods to do reading-related
-
- operations.
- """
+ """Abstract class that exposes methods to do reading-related operations."""
@abc.abstractmethod
def __init__(self):
@@ -227,10 +225,7 @@
class TypeTreeNode(object):
- """TypeTreeNode class exposes methods to format object tree structure
-
- data.
- """
+ """TypeTreeNode class exposes methods to format object tree structure data."""
def __init__(self, name, type_id, fields=None, permutation=None):
"""Constructs a new instance of TypeTreeNode.
@@ -692,14 +687,14 @@
class IgniteDataset(dataset_ops.DatasetSource):
- """Apache Ignite is a memory-centric distributed database, caching, and
+ """Apache Ignite is a memory-centric distributed database.
- processing platform for transactional, analytical, and streaming workloads,
- delivering in-memory speeds at petabyte scale. This contrib package
- contains an integration between Apache Ignite and TensorFlow. The
- integration is based on tf.data from TensorFlow side and Binary Client
- Protocol from Apache Ignite side. It allows to use Apache Ignite as a
- datasource for neural network training, inference and all other
+ It acts as a caching and processing platform for transactional, analytical,
+ and streaming workloads, delivering in-memory speeds at petabyte scale.
+ This contrib package contains an integration between Apache Ignite and
+ TensorFlow. The integration is based on tf.data from TensorFlow side and
+ Binary Client Protocol from Apache Ignite side. It allows to use Apache
+ Ignite as a datasource for neural network training, inference and all other
computations supported by TensorFlow. Ignite Dataset is based on Apache
Ignite Binary Client Protocol.
"""
@@ -756,6 +751,9 @@
self.cache_type.to_permutation(),
dtype=dtypes.int32,
name="permutation")
+ self._structure = structure.convert_legacy_structure(
+ self.cache_type.to_output_types(), self.cache_type.to_output_shapes(),
+ self.cache_type.to_output_classes())
def _as_variant_tensor(self):
return gen_dataset_ops.ignite_dataset(self.cache_name, self.host, self.port,
@@ -763,13 +761,5 @@
self.schema, self.permutation)
@property
- def output_classes(self):
- return self.cache_type.to_output_classes()
-
- @property
- def output_shapes(self):
- return self.cache_type.to_output_shapes()
-
- @property
- def output_types(self):
- return self.cache_type.to_output_types()
+ def _element_structure(self):
+ return self._structure
diff --git a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
index 4997c31..ba5cdfe 100644
--- a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
+++ b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
@@ -281,6 +281,13 @@
value.eval(),
np.array([[4, 4], [4, 4]]).astype(dtype.as_numpy_dtype()))
+ @test_util.run_in_graph_and_eager_modes
+ def test_transform_eager(self):
+ image = constant_op.constant([[1., 2.], [3., 4.]])
+ value = image_ops.transform(image, [1] * 8)
+ with self.test_session(use_gpu=True):
+ self.assertAllEqual(self.evaluate(value), np.array([[4, 4], [4, 4]]))
+
class BipartiteMatchTest(test_util.TensorFlowTestCase):
diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py
index d4fb99a..b25a6f7 100644
--- a/tensorflow/contrib/image/python/ops/image_ops.py
+++ b/tensorflow/contrib/image/python/ops/image_ops.py
@@ -17,6 +17,7 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.eager import context
from tensorflow.contrib.image.ops import gen_image_ops
from tensorflow.contrib.util import loader
from tensorflow.python.framework import common_shapes
@@ -271,8 +272,11 @@
raise TypeError("Images should have rank between 2 and 4.")
if output_shape is None:
- output_shape = tensor_util.constant_value(
- array_ops.shape(images)[1:3]) or array_ops.shape(images)[1:3]
+ output_shape = array_ops.shape(images)[1:3]
+ if not context.executing_eagerly():
+ output_shape_value = tensor_util.constant_value(output_shape)
+ if output_shape_value is not None:
+ output_shape = output_shape_value
output_shape = ops.convert_to_tensor(
output_shape, dtypes.int32, name="output_shape")
diff --git a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py
index 7129f09..2b86331 100644
--- a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py
+++ b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py
@@ -20,9 +20,9 @@
from tensorflow.contrib.kafka.python.ops import gen_dataset_ops
from tensorflow.contrib.kafka.python.ops import kafka_op_loader # pylint: disable=unused-import
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
class KafkaDataset(dataset_ops.DatasetSource):
@@ -63,13 +63,5 @@
self._group, self._eof, self._timeout)
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.string
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
diff --git a/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py b/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py
index c392adb..2039539 100644
--- a/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py
+++ b/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py
@@ -20,9 +20,9 @@
from tensorflow.contrib.kinesis.python.ops import gen_dataset_ops
from tensorflow.contrib.kinesis.python.ops import kinesis_op_loader # pylint: disable=unused-import
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
class KinesisDataset(dataset_ops.DatasetSource):
@@ -81,13 +81,5 @@
self._stream, self._shard, self._read_indefinitely, self._interval)
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.string
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py
index 0a4d2c6..d791418 100644
--- a/tensorflow/contrib/layers/python/layers/layers_test.py
+++ b/tensorflow/contrib/layers/python/layers/layers_test.py
@@ -1459,13 +1459,6 @@
class FlattenTest(test.TestCase):
- def testInvalidRank(self):
- with ops.Graph().as_default() as g, self.session(g):
- inputs = array_ops.placeholder(dtype=dtypes.float32)
- inputs.set_shape(tensor_shape.TensorShape((5,)))
- with self.assertRaisesRegexp(ValueError, 'incompatible with the layer'):
- _layers.flatten(inputs)
-
def testUnknownLastDim(self):
with ops.Graph().as_default() as g, self.session(g):
inputs = array_ops.placeholder(dtype=dtypes.float32)
@@ -1502,6 +1495,12 @@
images.get_shape().num_elements())
self.assertEqual(output.get_shape()[0], images.get_shape()[0])
+ def testFlatten0D(self):
+ with self.cached_session():
+ scalars = random_ops.random_uniform((5,), seed=1, name='scalars')
+ output = _layers.flatten(scalars)
+ self.assertEqual(output.shape, (5, 1))
+
def testFlattenBatchSize(self):
height, width = 3, 3
with self.cached_session() as sess:
diff --git a/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py b/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py
index 8466dc3..d49834d 100644
--- a/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py
+++ b/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py
@@ -12,7 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for SdcaModel."""
+"""Tests for SdcaModel (deprecated).
+
+This module and all its submodules are deprecated. To UPDATE or USE linear
+optimizers, please check its latest version in core:
+tensorflow_estimator/python/estimator/canned/linear_optimizer/.
+"""
from __future__ import absolute_import
from __future__ import division
diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py b/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py
index f3f1dcd9..c056a12 100644
--- a/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py
+++ b/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py
@@ -12,7 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Proximal stochastic dual coordinate ascent optimizer for linear models."""
+# pylint: disable=line-too-long
+"""Proximal stochastic dual coordinate ascent optimizer for linear models (deprecated).
+
+This module and all its submodules are deprecated. To UPDATE or USE linear
+optimizers, please check its latest version in core:
+tensorflow_estimator/python/estimator/canned/linear_optimizer/.
+"""
+# pylint: enable=line-too-long
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -40,6 +47,7 @@
from tensorflow.python.ops.nn import log_poisson_loss
from tensorflow.python.ops.nn import sigmoid_cross_entropy_with_logits
from tensorflow.python.summary import summary
+from tensorflow.python.util import deprecation
__all__ = ['SdcaModel']
@@ -48,7 +56,7 @@
class SdcaModel(object):
"""Stochastic dual coordinate ascent solver for linear models.
- Loss functions supported:
+ Loss functions supported:
* Binary logistic loss
* Squared loss
@@ -109,6 +117,10 @@
```
"""
+ @deprecation.deprecated(
+ None, 'This class is deprecated. To UPDATE or USE linear optimizers, '
+ 'please check its latest version in core: '
+ 'tensorflow_estimator/python/estimator/canned/linear_optimizer/.')
def __init__(self, examples, variables, options):
"""Create a new sdca optimizer."""
diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py
index a001555..a283949 100644
--- a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py
+++ b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py
@@ -12,7 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Sharded mutable dense hash table."""
+"""Sharded mutable dense hash table (deprecated).
+
+This module and all its submodules are deprecated. To UPDATE or USE linear
+optimizers, please check its latest version in core:
+tensorflow_estimator/python/estimator/canned/linear_optimizer/.
+"""
from __future__ import absolute_import
from __future__ import division
@@ -28,6 +33,7 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
+from tensorflow.python.util import deprecation
# TODO(rohanj): This should subclass Checkpointable and implement
@@ -45,6 +51,10 @@
# TODO(andreasst): consider moving this to lookup module
+ @deprecation.deprecated(
+ None, 'This class is deprecated. To UPDATE or USE linear optimizers, '
+ 'please check its latest version in core: '
+ 'tensorflow_estimator/python/estimator/canned/linear_optimizer/.')
def __init__(self,
key_dtype,
value_dtype,
diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py
index 2b56d0f..2d1457f 100644
--- a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py
+++ b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py
@@ -12,7 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for sharded_mutable_dense_hashtable.py."""
+"""Tests for sharded_mutable_dense_hashtable.py (deprecated).
+
+This module and all its submodules are deprecated. To UPDATE or USE linear
+optimizers, please check its latest version in core:
+tensorflow_estimator/python/estimator/canned/linear_optimizer/.
+"""
from __future__ import absolute_import
from __future__ import division
diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py
index 0037952..64730f8 100644
--- a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py
+++ b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py
@@ -12,7 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Sparse feature column."""
+"""Sparse feature column (deprecated).
+
+This module and all its submodules are deprecated. To UPDATE or USE linear
+optimizers, please check its latest version in core:
+tensorflow_estimator/python/estimator/canned/linear_optimizer/.
+"""
from __future__ import absolute_import
from __future__ import division
@@ -21,6 +26,7 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework.ops import internal_convert_to_tensor
from tensorflow.python.framework.ops import name_scope
+from tensorflow.python.util import deprecation
class SparseFeatureColumn(object):
@@ -68,6 +74,10 @@
@@feature_values
"""
+ @deprecation.deprecated(
+ None, 'This class is deprecated. To UPDATE or USE linear optimizers, '
+ 'please check its latest version in core: '
+ 'tensorflow_estimator/python/estimator/canned/linear_optimizer/.')
def __init__(self, example_indices, feature_indices, feature_values):
"""Creates a `SparseFeatureColumn` representation.
diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py
index 51c4f68..0ae780e 100644
--- a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py
+++ b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py
@@ -12,7 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for sparse_feature_column.py."""
+"""Tests for sparse_feature_column.py (deprecated).
+
+This module and all its submodules are deprecated. To UPDATE or USE linear
+optimizers, please check its latest version in core:
+tensorflow_estimator/python/estimator/canned/linear_optimizer/.
+"""
from __future__ import absolute_import
from __future__ import division
diff --git a/tensorflow/contrib/makefile/download_dependencies.sh b/tensorflow/contrib/makefile/download_dependencies.sh
index b396c52..2a5232b 100755
--- a/tensorflow/contrib/makefile/download_dependencies.sh
+++ b/tensorflow/contrib/makefile/download_dependencies.sh
@@ -30,11 +30,13 @@
GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)"
GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz"
NSYNC_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/nsync/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)"
-# Note: The Protobuf source in `tensorflow/workspace.bzl` in TensorFlow
-# 1.10 branch does not work. `make distclean` fails and blocks the build
-# process. For now we're hardcoding to the version which is used by
-# TensorFlow 1.9.
-PROTOBUF_URL="https://mirror.bazel.build/github.com/google/protobuf/archive/v3.6.0.tar.gz"
+
+# Note: The protobuf repo needs to be cloned due to its submodules.
+# These variables contain the GitHub repo and the sha, from `tensorflow/workspace.bzl`,
+# from which to clone it from and checkout to.
+readonly PROTOBUF_REPO="https://github.com/protocolbuffers/protobuf.git"
+readonly PROTOBUF_TAG="$(grep -o 'https://github.com/protocolbuffers/protobuf/archive/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1 | awk '{print substr($0, index($0, "archive") + 8, index($0, "tar") - index($0, "archive") - 9) }')"
+
# TODO (yongtang): Replace the following with 'https://mirror.bazel.build/github.com/google/re2/.*tar\.gz' once
# the archive has been propagated in mirror.bazel.build.
RE2_URL="$(grep -o 'https://github.com/google/re2/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)"
@@ -91,11 +93,34 @@
find "${dir}" -type f -name '*BUILD' -delete
}
+function clone_repository() {
+ local repo_url="${1}"
+ local destination_directory="${2}"
+ local commit_sha="${3}"
+
+ if [[ -d "${destination_directory}" ]]; then
+ rm -rf "${destination_directory}"
+ fi
+
+ git clone "${repo_url}" "${destination_directory}"
+
+ pushd "$(pwd)" 1>/dev/null
+
+ cd "${destination_directory}"
+
+ if [[ -n "${commit_sha}" ]]; then
+ git checkout "${PROTOBUF_TAG}"
+ fi
+
+ git submodule update --init
+
+ popd 1>/dev/null
+}
+
download_and_extract "${EIGEN_URL}" "${DOWNLOADS_DIR}/eigen"
download_and_extract "${GEMMLOWP_URL}" "${DOWNLOADS_DIR}/gemmlowp"
download_and_extract "${GOOGLETEST_URL}" "${DOWNLOADS_DIR}/googletest"
download_and_extract "${NSYNC_URL}" "${DOWNLOADS_DIR}/nsync"
-download_and_extract "${PROTOBUF_URL}" "${DOWNLOADS_DIR}/protobuf"
download_and_extract "${RE2_URL}" "${DOWNLOADS_DIR}/re2"
download_and_extract "${FFT2D_URL}" "${DOWNLOADS_DIR}/fft2d"
download_and_extract "${DOUBLE_CONVERSION_URL}" "${DOWNLOADS_DIR}/double_conversion"
@@ -106,6 +131,8 @@
download_and_extract "${FARMHASH_URL}" "${DOWNLOADS_DIR}/farmhash"
download_and_extract "${FLATBUFFERS_URL}" "${DOWNLOADS_DIR}/flatbuffers"
+clone_repository "${PROTOBUF_REPO}" "${DOWNLOADS_DIR}/protobuf" "${PROTOBUF_TAG}"
+
replace_by_sed 's#static uint32x4_t p4ui_CONJ_XOR = vld1q_u32( conj_XOR_DATA );#static uint32x4_t p4ui_CONJ_XOR; // = vld1q_u32( conj_XOR_DATA ); - Removed by script#' \
"${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h"
replace_by_sed 's#static uint32x2_t p2ui_CONJ_XOR = vld1_u32( conj_XOR_DATA );#static uint32x2_t p2ui_CONJ_XOR;// = vld1_u32( conj_XOR_DATA ); - Removed by scripts#' \
diff --git a/tensorflow/contrib/metrics/python/metrics/classification.py b/tensorflow/contrib/metrics/python/metrics/classification.py
index 062deb7..9aabc4b 100644
--- a/tensorflow/contrib/metrics/python/metrics/classification.py
+++ b/tensorflow/contrib/metrics/python/metrics/classification.py
@@ -18,13 +18,13 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import metrics_impl
from tensorflow.python.ops import variable_scope
-from tensorflow.python.training import distribution_strategy_context
# TODO(nsilberman): move into metrics/python/ops/
diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD
index f4ac70e..0446e82 100644
--- a/tensorflow/contrib/opt/BUILD
+++ b/tensorflow/contrib/opt/BUILD
@@ -14,6 +14,7 @@
name = "opt_py",
srcs = [
"__init__.py",
+ "python/training/adam_gs_optimizer.py",
"python/training/adamax.py",
"python/training/addsign.py",
"python/training/agn_optimizer.py",
@@ -22,6 +23,7 @@
"python/training/external_optimizer.py",
"python/training/ggt.py",
"python/training/lars_optimizer.py",
+ "python/training/lazy_adam_gs_optimizer.py",
"python/training/lazy_adam_optimizer.py",
"python/training/matrix_functions.py",
"python/training/model_average_optimizer.py",
@@ -61,6 +63,21 @@
)
py_test(
+ name = "adam_gs_optimizer_test",
+ srcs = ["python/training/adam_gs_optimizer_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":opt_py",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:framework_for_generated_wrappers",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:training",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
name = "adamax_test",
srcs = ["python/training/adamax_test.py"],
srcs_version = "PY2AND3",
@@ -149,6 +166,25 @@
)
py_test(
+ name = "lazy_adam_gs_optimizer_test",
+ srcs = ["python/training/lazy_adam_gs_optimizer_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":opt_py",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:resource_variable_ops",
+ "//tensorflow/python:variables",
+ "//third_party/py/numpy",
+ "@absl_py//absl/testing:parameterized",
+ ],
+)
+
+py_test(
name = "lazy_adam_optimizer_test",
srcs = ["python/training/lazy_adam_optimizer_test.py"],
srcs_version = "PY2AND3",
diff --git a/tensorflow/contrib/opt/__init__.py b/tensorflow/contrib/opt/__init__.py
index c7ea68e..e8fc523 100644
--- a/tensorflow/contrib/opt/__init__.py
+++ b/tensorflow/contrib/opt/__init__.py
@@ -19,6 +19,7 @@
from __future__ import print_function
# pylint: disable=wildcard-import
+from tensorflow.contrib.opt.python.training.adam_gs_optimizer import *
from tensorflow.contrib.opt.python.training.adamax import *
from tensorflow.contrib.opt.python.training.addsign import *
from tensorflow.contrib.opt.python.training.agn_optimizer import *
@@ -28,6 +29,7 @@
from tensorflow.contrib.opt.python.training.lars_optimizer import *
from tensorflow.contrib.opt.python.training.ggt import *
from tensorflow.contrib.opt.python.training.lazy_adam_optimizer import *
+from tensorflow.contrib.opt.python.training.lazy_adam_gs_optimizer import *
from tensorflow.contrib.opt.python.training.model_average_optimizer import *
from tensorflow.contrib.opt.python.training.moving_average_optimizer import *
from tensorflow.contrib.opt.python.training.multitask_optimizer_wrapper import *
@@ -44,12 +46,14 @@
_allowed_symbols = [
'AdaMaxOptimizer',
+ 'AdamGSOptimizer',
'PowerSignOptimizer',
'AddSignOptimizer',
'DelayCompensatedGradientDescentOptimizer',
'DropStaleGradientOptimizer',
'ExternalOptimizerInterface',
'LARSOptimizer',
+ 'LazyAdamGSOptimizer',
'LazyAdamOptimizer',
'NadamOptimizer',
'MovingAverageOptimizer',
diff --git a/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py b/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py
new file mode 100644
index 0000000..3fb649e
--- /dev/null
+++ b/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py
@@ -0,0 +1,217 @@
+# 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.
+# ==============================================================================
+
+"""Adam rewrite to use global step for computing beta1 & beta2 accumulation."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.framework import dtypes
+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 math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("train.AdamOptimizer")
+class AdamGSOptimizer(optimizer.Optimizer):
+ """Optimizer that implements the Adam algorithm.
+
+ See [Kingma et al., 2014](http://arxiv.org/abs/1412.6980)
+ ([pdf](http://arxiv.org/pdf/1412.6980.pdf)).
+ """
+
+ def __init__(self, global_step=0, learning_rate=0.001,
+ beta1=0.9, beta2=0.999, epsilon=1e-8,
+ use_locking=False, name="Adam"):
+ """Construct a new Adam optimizer.
+
+ Branched from tf.train.AdamOptimizer. The only difference is to pass
+ global step for computing beta1 and beta2 accumulators, instead of having
+ optimizer keep its own independent beta1 and beta2 accumulators as non-slot
+ variables.
+
+ Initialization:
+
+ $$m_0 := 0 \text{(Initialize initial 1st moment vector)}$$
+ $$v_0 := 0 \text{(Initialize initial 2nd moment vector)}$$
+ $$t := 0 \text{(Initialize timestep)}$$
+
+ The update rule for `variable` with gradient `g` uses an optimization
+ described at the end of section2 of the paper:
+
+ $$t := t + 1$$
+ $$lr_t := \text{learning\_rate} * \sqrt{1 - beta_2^t} / (1 - beta_1^t)$$
+
+ $$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$
+ $$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$
+ $$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$
+
+ The default value of 1e-8 for epsilon might not be a good default in
+ general. For example, when training an Inception network on ImageNet a
+ current good choice is 1.0 or 0.1. Note that since AdamOptimizer uses the
+ formulation just before Section 2.1 of the Kingma and Ba paper rather than
+ the formulation in Algorithm 1, the "epsilon" referred to here is "epsilon
+ hat" in the paper.
+
+ The sparse implementation of this algorithm (used when the gradient is an
+ IndexedSlices object, typically because of `tf.gather` or an embedding
+ lookup in the forward pass) does apply momentum to variable slices even if
+ they were not used in the forward pass (meaning they have a gradient equal
+ to zero). Momentum decay (beta1) is also applied to the entire momentum
+ accumulator. This means that the sparse behavior is equivalent to the dense
+ behavior (in contrast to some momentum implementations which ignore momentum
+ unless a variable slice was actually used).
+
+ Args:
+ global_step: tensorflow variable indicating the step.
+ learning_rate: A Tensor or a floating point value. The learning rate.
+ beta1: A float value or a constant float tensor.
+ The exponential decay rate for the 1st moment estimates.
+ beta2: A float value or a constant float tensor.
+ The exponential decay rate for the 2nd moment estimates.
+ epsilon: A small constant for numerical stability. This epsilon is
+ "epsilon hat" in the Kingma and Ba paper (in the formula just before
+ Section 2.1), not the epsilon in Algorithm 1 of the paper.
+ use_locking: If True use locks for update operations.
+ name: Optional name for the operations created when applying gradients.
+ Defaults to "Adam".
+
+ @compatibility(eager)
+ When eager execution is enabled, `learning_rate`, `beta1`, `beta2`, and
+ `epsilon` can each be a callable that takes no arguments and returns the
+ actual value to use. This can be useful for changing these values across
+ different invocations of optimizer functions.
+ @end_compatibility
+ """
+ super(AdamGSOptimizer, self).__init__(use_locking, name)
+ self._lr = learning_rate
+ self._beta1 = beta1
+ self._beta2 = beta2
+ self._epsilon = epsilon
+ self._global_step = global_step
+ self._global_step_on_worker = None
+
+ # Tensor versions of the constructor arguments, created in _prepare().
+ self._lr_t = None
+ self._beta1_t = None
+ self._beta2_t = None
+ self._epsilon_t = None
+
+ # Created in SparseApply if needed.
+ self._updated_lr = None
+
+ def _get_beta_accumulators(self):
+ return (math_ops.pow(self._beta1_t, self._global_step_on_worker),
+ math_ops.pow(self._beta2_t, self._global_step_on_worker))
+
+ def _create_slots(self, var_list):
+ # Create slots for the first and second moments.
+ for v in var_list:
+ self._zeros_slot(v, "m", self._name)
+ self._zeros_slot(v, "v", self._name)
+
+ def _prepare(self):
+ lr = self._call_if_callable(self._lr)
+ beta1 = self._call_if_callable(self._beta1)
+ beta2 = self._call_if_callable(self._beta2)
+ epsilon = self._call_if_callable(self._epsilon)
+
+ self._lr_t = ops.convert_to_tensor(lr, name="learning_rate")
+ self._beta1_t = ops.convert_to_tensor(beta1, name="beta1")
+ self._beta2_t = ops.convert_to_tensor(beta2, name="beta2")
+ self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon")
+
+ # Performance optimization so that worker creates a copy of the global step
+ # to avoid overloading the parameter server holding the global step.
+ self._global_step_on_worker = math_ops.cast(
+ array_ops.identity(self._global_step) + 1, dtypes.float32)
+
+ def _apply_dense(self, grad, var):
+ m = self.get_slot(var, "m")
+ v = self.get_slot(var, "v")
+ beta1_power, beta2_power = self._get_beta_accumulators()
+ return training_ops.apply_adam(
+ var, m, v,
+ math_ops.cast(beta1_power, var.dtype.base_dtype),
+ math_ops.cast(beta2_power, var.dtype.base_dtype),
+ math_ops.cast(self._lr_t, var.dtype.base_dtype),
+ math_ops.cast(self._beta1_t, var.dtype.base_dtype),
+ math_ops.cast(self._beta2_t, var.dtype.base_dtype),
+ math_ops.cast(self._epsilon_t, var.dtype.base_dtype),
+ grad, use_locking=self._use_locking).op
+
+ def _resource_apply_dense(self, grad, var):
+ m = self.get_slot(var, "m")
+ v = self.get_slot(var, "v")
+ beta1_power, beta2_power = self._get_beta_accumulators()
+ return training_ops.resource_apply_adam(
+ var.handle, m.handle, v.handle,
+ math_ops.cast(beta1_power, grad.dtype.base_dtype),
+ math_ops.cast(beta2_power, grad.dtype.base_dtype),
+ math_ops.cast(self._lr_t, grad.dtype.base_dtype),
+ math_ops.cast(self._beta1_t, grad.dtype.base_dtype),
+ math_ops.cast(self._beta2_t, grad.dtype.base_dtype),
+ math_ops.cast(self._epsilon_t, grad.dtype.base_dtype),
+ grad, use_locking=self._use_locking)
+
+ def _apply_sparse_shared(self, grad, var, indices, scatter_add):
+ beta1_power, beta2_power = self._get_beta_accumulators()
+ beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
+ beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
+ lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
+ beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
+ beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
+ epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
+ lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
+ # m_t = beta1 * m + (1 - beta1) * g_t
+ m = self.get_slot(var, "m")
+ m_scaled_g_values = grad * (1 - beta1_t)
+ m_t = state_ops.assign(m, m * beta1_t,
+ use_locking=self._use_locking)
+ with ops.control_dependencies([m_t]):
+ m_t = scatter_add(m, indices, m_scaled_g_values)
+ # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+ v = self.get_slot(var, "v")
+ v_scaled_g_values = (grad * grad) * (1 - beta2_t)
+ v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking)
+ with ops.control_dependencies([v_t]):
+ v_t = scatter_add(v, indices, v_scaled_g_values)
+ v_sqrt = math_ops.sqrt(v_t)
+ var_update = state_ops.assign_sub(var,
+ lr * m_t / (v_sqrt + epsilon_t),
+ use_locking=self._use_locking)
+ return control_flow_ops.group(*[var_update, m_t, v_t])
+
+ def _apply_sparse(self, grad, var):
+ return self._apply_sparse_shared(
+ grad.values, var, grad.indices,
+ lambda x, i, v: state_ops.scatter_add( # pylint: disable=g-long-lambda
+ x, i, v, use_locking=self._use_locking))
+
+ def _resource_scatter_add(self, x, i, v):
+ with ops.control_dependencies(
+ [resource_variable_ops.resource_scatter_add(
+ x.handle, i, v)]):
+ return x.value()
+
+ def _resource_apply_sparse(self, grad, var, indices):
+ return self._apply_sparse_shared(
+ grad, var, indices, self._resource_scatter_add)
diff --git a/tensorflow/contrib/opt/python/training/adam_gs_optimizer_test.py b/tensorflow/contrib/opt/python/training/adam_gs_optimizer_test.py
new file mode 100644
index 0000000..c68c965
--- /dev/null
+++ b/tensorflow/contrib/opt/python/training/adam_gs_optimizer_test.py
@@ -0,0 +1,382 @@
+# 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 AdamGS."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.contrib.opt.python.training import adam_gs_optimizer
+from tensorflow.python.client import session
+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.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.platform import test
+
+
+def adam_update_numpy(param,
+ g_t,
+ t,
+ m,
+ v,
+ alpha=0.001,
+ beta1=0.9,
+ beta2=0.999,
+ epsilon=1e-8):
+ alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t)
+
+ m_t = beta1 * m + (1 - beta1) * g_t
+ v_t = beta2 * v + (1 - beta2) * g_t * g_t
+
+ param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon)
+ return param_t, m_t, v_t
+
+
+class AdamGSOptimizerTest(test.TestCase):
+
+ def doTestSparse(self, use_resource=False):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ if use_resource:
+ global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64))
+ var0 = resource_variable_ops.ResourceVariable(var0_np)
+ var1 = resource_variable_ops.ResourceVariable(var1_np)
+ else:
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0_np_indices = np.array([0, 1], dtype=np.int32)
+ grads0 = ops.IndexedSlices(
+ constant_op.constant(grads0_np),
+ constant_op.constant(grads0_np_indices), constant_op.constant([2]))
+ grads1_np_indices = np.array([0, 1], dtype=np.int32)
+ grads1 = ops.IndexedSlices(
+ constant_op.constant(grads1_np),
+ constant_op.constant(grads1_np_indices), constant_op.constant([2]))
+ opt = adam_gs_optimizer.AdamGSOptimizer(global_step=global_step)
+ update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ variables.global_variables_initializer().run()
+
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+ self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+
+ # Run 3 steps of Adam
+ for t in range(1, 4):
+ self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(0.999**t,
+ self.evaluate(beta2_power))
+ update.run()
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+ self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+ def testSparse(self):
+ self.doTestSparse(use_resource=False)
+
+ def testResourceSparse(self):
+ self.doTestSparse(use_resource=True)
+
+ def testSparseDevicePlacement(self):
+ for index_dtype in [dtypes.int32, dtypes.int64]:
+ with self.cached_session(force_gpu=test.is_gpu_available()):
+ # If a GPU is available, tests that all optimizer ops can be placed on
+ # it (i.e. they have GPU kernels).
+ var = variables.Variable([[1.0], [2.0]])
+ indices = constant_op.constant([0, 1], dtype=index_dtype)
+ gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices))
+ optimizer = adam_gs_optimizer.AdamGSOptimizer(3.0)
+ minimize_op = optimizer.minimize(gathered_sum)
+ variables.global_variables_initializer().run()
+ minimize_op.run()
+
+ def testSparseRepeatedIndices(self):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ repeated_index_global_step = variables.Variable(
+ array_ops.zeros([], dtypes.int64))
+ aggregated_global_step = variables.Variable(
+ array_ops.zeros([], dtypes.int64))
+ repeated_index_update_var = variables.Variable(
+ [[1.0], [2.0]], dtype=dtype)
+ aggregated_update_var = variables.Variable(
+ [[1.0], [2.0]], dtype=dtype)
+ grad_repeated_index = ops.IndexedSlices(
+ constant_op.constant(
+ [0.1, 0.1], shape=[2, 1], dtype=dtype),
+ constant_op.constant([1, 1]),
+ constant_op.constant([2, 1]))
+ grad_aggregated = ops.IndexedSlices(
+ constant_op.constant(
+ [0.2], shape=[1, 1], dtype=dtype),
+ constant_op.constant([1]),
+ constant_op.constant([2, 1]))
+ repeated_update = adam_gs_optimizer.AdamGSOptimizer(
+ global_step=repeated_index_global_step).apply_gradients(
+ [(grad_repeated_index, repeated_index_update_var)],
+ global_step=repeated_index_global_step)
+ aggregated_update = adam_gs_optimizer.AdamGSOptimizer(
+ global_step=aggregated_global_step).apply_gradients(
+ [(grad_aggregated, aggregated_update_var)],
+ global_step=aggregated_global_step)
+ variables.global_variables_initializer().run()
+ self.assertAllClose(aggregated_update_var.eval(),
+ self.evaluate(repeated_index_update_var))
+ for _ in range(3):
+ repeated_update.run()
+ aggregated_update.run()
+ self.assertAllClose(aggregated_update_var.eval(),
+ self.evaluate(repeated_index_update_var))
+
+ def doTestBasic(self, use_resource=False, use_callable_params=False):
+ for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+ with self.session(graph=ops.Graph()):
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ if use_resource:
+ global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64), name="global_step_%d" % i)
+ var0 = resource_variable_ops.ResourceVariable(
+ var0_np, name="var0_%d" % i)
+ var1 = resource_variable_ops.ResourceVariable(
+ var1_np, name="var1_%d" % i)
+ else:
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0 = constant_op.constant(grads0_np)
+ grads1 = constant_op.constant(grads1_np)
+
+ learning_rate = lambda: 0.001
+ beta1 = lambda: 0.9
+ beta2 = lambda: 0.999
+ epsilon = lambda: 1e-8
+ if not use_callable_params:
+ learning_rate = learning_rate()
+ beta1 = beta1()
+ beta2 = beta2()
+ epsilon = epsilon()
+
+ opt = adam_gs_optimizer.AdamGSOptimizer(global_step=global_step,
+ learning_rate=learning_rate)
+ update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ opt_variables = opt.variables()
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+ self.assertTrue(beta1_power is not None)
+ self.assertTrue(beta2_power is not None)
+ self.assertNotIn(beta1_power, opt_variables)
+ self.assertNotIn(beta2_power, opt_variables)
+
+ if not context.executing_eagerly():
+ with ops.Graph().as_default():
+ # Shouldn't return non-slot variables from other graphs.
+ self.assertEqual(0, len(opt.variables()))
+ self.evaluate(variables.global_variables_initializer())
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+ self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+ # Run 3 steps of Adam
+ for t in range(1, 4):
+ if not context.executing_eagerly():
+ self.evaluate(update)
+ self.assertAllCloseAccordingToType(
+ 0.9**(t + 1), self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(
+ 0.999**(t + 1), self.evaluate(beta2_power))
+ else:
+ if t > 1:
+ opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+ self.assertAllCloseAccordingToType(
+ 0.9**t, self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(
+ 0.999**t, self.evaluate(beta2_power))
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+ self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+ if use_resource:
+ self.assertEqual("var0_%d/Adam:0" % (i,),
+ opt.get_slot(var=var0, name="m").name)
+
+ def testBasic(self):
+ with self.cached_session():
+ self.doTestBasic(use_resource=False)
+
+ @test_util.run_in_graph_and_eager_modes(reset_test=True)
+ def testResourceBasic(self):
+ self.doTestBasic(use_resource=True)
+
+ def testBasicCallableParams(self):
+ with context.eager_mode():
+ self.doTestBasic(use_resource=True, use_callable_params=True)
+
+ def testTensorLearningRate(self):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0 = constant_op.constant(grads0_np)
+ grads1 = constant_op.constant(grads1_np)
+ opt = adam_gs_optimizer.AdamGSOptimizer(
+ global_step=global_step, learning_rate=constant_op.constant(0.001))
+ update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ variables.global_variables_initializer().run()
+
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+ self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+
+ # Run 3 steps of Adam
+ for t in range(1, 4):
+ self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(0.999**t,
+ self.evaluate(beta2_power))
+ update.run()
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+ self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+ def testSharing(self):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0 = constant_op.constant(grads0_np)
+ grads1 = constant_op.constant(grads1_np)
+ opt = adam_gs_optimizer.AdamGSOptimizer(global_step=global_step)
+ update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ variables.global_variables_initializer().run()
+
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+ self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+ # Run 3 steps of intertwined Adam1 and Adam2.
+ for t in range(1, 4):
+ self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(0.999**t,
+ self.evaluate(beta2_power))
+ if t % 2 == 0:
+ update1.run()
+ else:
+ update2.run()
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+ self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+ def testTwoSessions(self):
+ optimizer = adam_gs_optimizer.AdamGSOptimizer()
+
+ with context.eager_mode():
+ var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+ grads0 = constant_op.constant(np.array([0.1, 0.1]))
+ optimizer.apply_gradients([(grads0, var0)])
+
+ g = ops.Graph()
+ with g.as_default():
+ with session.Session():
+ var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+ grads0 = constant_op.constant(np.array([0.1, 0.1]))
+ optimizer.apply_gradients([(grads0, var0)])
+
+ gg = ops.Graph()
+ with gg.as_default():
+ with session.Session():
+ var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+ grads0 = constant_op.constant(np.array([0.1, 0.1]))
+
+ # If the optimizer saves any state not keyed by graph the following line
+ # fails.
+ optimizer.apply_gradients([(grads0, var0)])
+
+ def testSlotsUniqueEager(self):
+ with context.eager_mode():
+ v1 = resource_variable_ops.ResourceVariable(1.)
+ v2 = resource_variable_ops.ResourceVariable(1.)
+ opt = adam_gs_optimizer.AdamGSOptimizer(1.)
+ opt.minimize(lambda: v1 + v2)
+ # There should be two unique slot variables for v1 and v2 respectively.
+ self.assertEqual(4, len(set(opt.variables())))
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py
index 6c203e5..fa1a7aa 100644
--- a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py
+++ b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py
@@ -30,6 +30,7 @@
from tensorflow.python.training import optimizer
from tensorflow.python.training import saver
from tensorflow.python.training import session_run_hook
+from tensorflow.python.training.saving import saveable_object_util
LOCAL_VARIABLE_NAME = 'local_center_variable'
GLOBAL_VARIABLE_NAME = 'global_center_variable'
@@ -424,7 +425,7 @@
if var_list is None:
var_list = variables.trainable_variables()
if not isinstance(var_list, dict):
- var_list = saver.BaseSaverBuilder.OpListToDict(var_list)
+ var_list = saveable_object_util.op_list_to_dict(var_list)
swapped_var_list = {}
for key, var in var_list.items():
@@ -464,4 +465,4 @@
def after_create_session(self, session, coord):
"""Run initialization ops"""
- session.run(self._variable_init_op)
\ No newline at end of file
+ session.run(self._variable_init_op)
diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer.py b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer.py
new file mode 100644
index 0000000..8827007
--- /dev/null
+++ b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer.py
@@ -0,0 +1,114 @@
+# 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.
+# ==============================================================================
+
+"""LazyAdam rewrite to use global step for computing beta1 & beta2 accumulation.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.contrib.opt.python.training import adam_gs_optimizer
+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 state_ops
+
+
+class LazyAdamGSOptimizer(adam_gs_optimizer.AdamGSOptimizer):
+ """Variant of the Adam optimizer that handles sparse updates more efficiently.
+
+ Branched from tf.contrib.opt.LazyAdamGSOptimizer. The only difference is to
+ pass global step for computing beta1 and beta2 accumulators, instead of having
+ optimizer keep its own independent beta1 and beta2 accumulators as non-slot
+ variables.
+
+ The original Adam algorithm maintains two moving-average accumulators for
+ each trainable variable; the accumulators are updated at every step.
+ This class provides lazier handling of gradient updates for sparse variables.
+ It only updates moving-average accumulators for sparse variable indices that
+ appear in the current batch, rather than updating the accumulators for all
+ indices. Compared with the original Adam optimizer, it can provide large
+ improvements in model training throughput for some applications. However, it
+ provides slightly different semantics than the original Adam algorithm, and
+ may lead to different empirical results.
+ """
+
+ def _apply_sparse(self, grad, var):
+ beta1_power, beta2_power = self._get_beta_accumulators()
+ beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
+ beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
+ lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
+ beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
+ beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
+ epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
+ lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
+
+ # \\(m := beta1 * m + (1 - beta1) * g_t\\)
+ m = self.get_slot(var, "m")
+ m_t = state_ops.scatter_update(m, grad.indices,
+ beta1_t * array_ops.gather(m, grad.indices) +
+ (1 - beta1_t) * grad.values,
+ use_locking=self._use_locking)
+
+ # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\)
+ v = self.get_slot(var, "v")
+ v_t = state_ops.scatter_update(v, grad.indices,
+ beta2_t * array_ops.gather(v, grad.indices) +
+ (1 - beta2_t) * math_ops.square(grad.values),
+ use_locking=self._use_locking)
+
+ # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\)
+ m_t_slice = array_ops.gather(m_t, grad.indices)
+ v_t_slice = array_ops.gather(v_t, grad.indices)
+ denominator_slice = math_ops.sqrt(v_t_slice) + epsilon_t
+ var_update = state_ops.scatter_sub(var, grad.indices,
+ lr * m_t_slice / denominator_slice,
+ use_locking=self._use_locking)
+ return control_flow_ops.group(var_update, m_t, v_t)
+
+ def _resource_apply_sparse(self, grad, var, indices):
+ beta1_power, beta2_power = self._get_beta_accumulators()
+ beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
+ beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
+ lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
+ beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
+ beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
+ epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
+ lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
+
+ # \\(m := beta1 * m + (1 - beta1) * g_t\\)
+ m = self.get_slot(var, "m")
+ m_t_slice = beta1_t * array_ops.gather(m, indices) + (1 - beta1_t) * grad
+ m_update_op = resource_variable_ops.resource_scatter_update(m.handle,
+ indices,
+ m_t_slice)
+
+ # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\)
+ v = self.get_slot(var, "v")
+ v_t_slice = (beta2_t * array_ops.gather(v, indices) +
+ (1 - beta2_t) * math_ops.square(grad))
+ v_update_op = resource_variable_ops.resource_scatter_update(v.handle,
+ indices,
+ v_t_slice)
+
+ # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\)
+ var_slice = lr * m_t_slice / (math_ops.sqrt(v_t_slice) + epsilon_t)
+ var_update_op = resource_variable_ops.resource_scatter_sub(var.handle,
+ indices,
+ var_slice)
+
+ return control_flow_ops.group(var_update_op, m_update_op, v_update_op)
diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer_test.py b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer_test.py
new file mode 100644
index 0000000..bdc9a02
--- /dev/null
+++ b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer_test.py
@@ -0,0 +1,402 @@
+# 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 LazyAdamGSOptimizer."""
+
+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.contrib.opt.python.training import lazy_adam_gs_optimizer
+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.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.platform import test
+
+
+def adam_update_numpy(param,
+ g_t,
+ t,
+ m,
+ v,
+ alpha=0.001,
+ beta1=0.9,
+ beta2=0.999,
+ epsilon=1e-8):
+ alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t)
+
+ m_t = beta1 * m + (1 - beta1) * g_t
+ v_t = beta2 * v + (1 - beta2) * g_t * g_t
+
+ param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon)
+ return param_t, m_t, v_t
+
+
+class LazyAdamGSOptimizerTest(test.TestCase, parameterized.TestCase):
+
+ @parameterized.parameters([False, True])
+ def testSparse(self, use_resource):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ if use_resource:
+ global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64))
+ var0 = resource_variable_ops.ResourceVariable(var0_np)
+ var1 = resource_variable_ops.ResourceVariable(var1_np)
+ else:
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+
+ grads0_np_indices = np.array([0, 1], dtype=np.int32)
+ grads0 = ops.IndexedSlices(
+ constant_op.constant(grads0_np),
+ constant_op.constant(grads0_np_indices), constant_op.constant([2]))
+ grads1_np_indices = np.array([0, 1], dtype=np.int32)
+ grads1 = ops.IndexedSlices(
+ constant_op.constant(grads1_np),
+ constant_op.constant(grads1_np_indices), constant_op.constant([2]))
+ opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=global_step)
+ update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ variables.global_variables_initializer().run()
+
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], var0.eval())
+ self.assertAllClose([3.0, 4.0], var1.eval())
+
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+
+ # Run 3 steps of Adam
+ for t in range(1, 4):
+ self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval())
+ self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval())
+ update.run()
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, var0.eval())
+ self.assertAllCloseAccordingToType(var1_np, var1.eval())
+
+ @parameterized.parameters([False, True])
+ def testSparseDevicePlacement(self, use_resource):
+ for index_dtype in [dtypes.int32, dtypes.int64]:
+ with self.cached_session(force_gpu=test.is_gpu_available()):
+ # If a GPU is available, tests that all optimizer ops can be placed on
+ # it (i.e. they have GPU kernels).
+ if use_resource:
+ global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64))
+ var = resource_variable_ops.ResourceVariable([[1.0], [2.0]])
+ else:
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ var = variables.Variable([[1.0], [2.0]])
+
+ indices = constant_op.constant([0, 1], dtype=index_dtype)
+ gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices))
+ optimizer = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=global_step, learning_rate=3.0)
+ minimize_op = optimizer.minimize(gathered_sum, global_step=global_step)
+ variables.global_variables_initializer().run()
+ minimize_op.run()
+
+ @parameterized.parameters([False, True])
+ def testSparseRepeatedIndices(self, use_resource):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ if use_resource:
+ repeated_index_global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64))
+ aggregated_global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64))
+ repeated_index_update_var = resource_variable_ops.ResourceVariable(
+ [[1.0], [2.0]], dtype=dtype)
+ aggregated_update_var = resource_variable_ops.ResourceVariable(
+ [[1.0], [2.0]], dtype=dtype)
+ else:
+ repeated_index_global_step = variables.Variable(
+ array_ops.zeros([], dtypes.int64))
+ aggregated_global_step = variables.Variable(
+ array_ops.zeros([], dtypes.int64))
+ repeated_index_update_var = variables.Variable(
+ [[1.0], [2.0]], dtype=dtype)
+ aggregated_update_var = variables.Variable(
+ [[1.0], [2.0]], dtype=dtype)
+
+ grad_repeated_index = ops.IndexedSlices(
+ constant_op.constant(
+ [0.1, 0.1], shape=[2, 1], dtype=dtype),
+ constant_op.constant([1, 1]),
+ constant_op.constant([2, 1]))
+ grad_aggregated = ops.IndexedSlices(
+ constant_op.constant(
+ [0.2], shape=[1, 1], dtype=dtype),
+ constant_op.constant([1]),
+ constant_op.constant([2, 1]))
+ repeated_update_opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=repeated_index_global_step)
+ repeated_update = repeated_update_opt.apply_gradients(
+ [(grad_repeated_index, repeated_index_update_var)],
+ global_step=repeated_index_global_step)
+ aggregated_update_opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=aggregated_global_step)
+ aggregated_update = aggregated_update_opt.apply_gradients(
+ [(grad_aggregated, aggregated_update_var)],
+ global_step=aggregated_global_step)
+ variables.global_variables_initializer().run()
+ self.assertAllClose(aggregated_update_var.eval(),
+ repeated_index_update_var.eval())
+ for _ in range(3):
+ repeated_update.run()
+ aggregated_update.run()
+ self.assertAllClose(aggregated_update_var.eval(),
+ repeated_index_update_var.eval())
+
+ def doTestBasic(self, use_resource=False, use_callable_params=False):
+ for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
+ with self.session(graph=ops.Graph()):
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ if use_resource:
+ global_step = resource_variable_ops.ResourceVariable(
+ array_ops.zeros([], dtypes.int64), name="global_step_%d" % i)
+ var0 = resource_variable_ops.ResourceVariable(
+ var0_np, name="var0_%d" % i)
+ var1 = resource_variable_ops.ResourceVariable(
+ var1_np, name="var1_%d" % i)
+ else:
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0 = constant_op.constant(grads0_np)
+ grads1 = constant_op.constant(grads1_np)
+
+ learning_rate = lambda: 0.001
+ beta1 = lambda: 0.9
+ beta2 = lambda: 0.999
+ epsilon = lambda: 1e-8
+ if not use_callable_params:
+ learning_rate = learning_rate()
+ beta1 = beta1()
+ beta2 = beta2()
+ epsilon = epsilon()
+
+ opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=global_step, learning_rate=learning_rate)
+ update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ opt_variables = opt.variables()
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+ self.assertIsNotNone(beta1_power)
+ self.assertIsNotNone(beta2_power is not None)
+ self.assertNotIn(beta1_power, opt_variables)
+ self.assertNotIn(beta2_power, opt_variables)
+
+ if not context.executing_eagerly():
+ with ops.Graph().as_default():
+ # Shouldn't return non-slot variables from other graphs.
+ self.assertEqual(0, len(opt.variables()))
+ self.evaluate(variables.global_variables_initializer())
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+ self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+ # Run 3 steps of Adam
+ for t in range(1, 4):
+ if not context.executing_eagerly():
+ self.evaluate(update)
+ self.assertAllCloseAccordingToType(
+ 0.9**(t + 1), self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(
+ 0.999**(t + 1), self.evaluate(beta2_power))
+ else:
+ if t > 1:
+ opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+ self.assertAllCloseAccordingToType(
+ 0.9**t, self.evaluate(beta1_power))
+ self.assertAllCloseAccordingToType(
+ 0.999**t, self.evaluate(beta2_power))
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+ self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+ if use_resource:
+ self.assertEqual("var0_%d/Adam:0" % (i,),
+ opt.get_slot(var=var0, name="m").name)
+
+ def testBasic(self):
+ with self.cached_session():
+ self.doTestBasic(use_resource=False)
+
+ @test_util.run_in_graph_and_eager_modes(reset_test=True)
+ def testResourceBasic(self):
+ self.doTestBasic(use_resource=True)
+
+ def testBasicCallableParams(self):
+ with context.eager_mode():
+ self.doTestBasic(use_resource=True, use_callable_params=True)
+
+ def testTensorLearningRate(self):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0 = constant_op.constant(grads0_np)
+ grads1 = constant_op.constant(grads1_np)
+ opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=global_step, learning_rate=constant_op.constant(0.001))
+ update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ variables.global_variables_initializer().run()
+
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], var0.eval())
+ self.assertAllClose([3.0, 4.0], var1.eval())
+
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+
+ # Run 3 steps of Adam
+ for t in range(1, 4):
+ self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval())
+ self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval())
+ update.run()
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, var0.eval())
+ self.assertAllCloseAccordingToType(var1_np, var1.eval())
+
+ def testSharing(self):
+ for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
+ with self.cached_session():
+ global_step = variables.Variable(array_ops.zeros([], dtypes.int64))
+ # Initialize variables for numpy implementation.
+ m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+ var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
+ grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
+ var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
+ grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
+
+ var0 = variables.Variable(var0_np)
+ var1 = variables.Variable(var1_np)
+ grads0 = constant_op.constant(grads0_np)
+ grads1 = constant_op.constant(grads1_np)
+ opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(
+ global_step=global_step)
+ update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]),
+ global_step=global_step)
+ variables.global_variables_initializer().run()
+
+ beta1_power, beta2_power = opt._get_beta_accumulators()
+
+ # Fetch params to validate initial values
+ self.assertAllClose([1.0, 2.0], var0.eval())
+ self.assertAllClose([3.0, 4.0], var1.eval())
+
+ # Run 3 steps of intertwined Adam1 and Adam2.
+ for t in range(1, 4):
+ self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval())
+ self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval())
+ if t % 2 == 0:
+ update1.run()
+ else:
+ update2.run()
+
+ var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+ var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+ # Validate updated params
+ self.assertAllCloseAccordingToType(var0_np, var0.eval())
+ self.assertAllCloseAccordingToType(var1_np, var1.eval())
+
+ def testTwoSessions(self):
+ optimizer = lazy_adam_gs_optimizer.LazyAdamGSOptimizer()
+
+ with context.eager_mode():
+ var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+ grads0 = constant_op.constant(np.array([0.1, 0.1]))
+ optimizer.apply_gradients([(grads0, var0)])
+
+ g = ops.Graph()
+ with g.as_default():
+ with self.session(graph=g):
+ var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+ grads0 = constant_op.constant(np.array([0.1, 0.1]))
+ optimizer.apply_gradients([(grads0, var0)])
+
+ gg = ops.Graph()
+ with gg.as_default():
+ with self.session(graph=gg):
+ var0 = variables.Variable(np.array([1.0, 2.0]), name="v0")
+ grads0 = constant_op.constant(np.array([0.1, 0.1]))
+
+ # If the optimizer saves any state not keyed by graph the following line
+ # fails.
+ optimizer.apply_gradients([(grads0, var0)])
+
+ def testSlotsUniqueEager(self):
+ with context.eager_mode():
+ v1 = resource_variable_ops.ResourceVariable(1.)
+ v2 = resource_variable_ops.ResourceVariable(1.)
+ opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(1.)
+ opt.minimize(lambda: v1 + v2)
+ # There should be two non-slot variables, and two unique slot variables
+ # for v1 and v2 respectively.
+ self.assertLen(set(opt.variables()), 4)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/contrib/opt/python/training/moving_average_optimizer.py b/tensorflow/contrib/opt/python/training/moving_average_optimizer.py
index b7fd2d2..bf3e5c5 100644
--- a/tensorflow/contrib/opt/python/training/moving_average_optimizer.py
+++ b/tensorflow/contrib/opt/python/training/moving_average_optimizer.py
@@ -26,6 +26,7 @@
from tensorflow.python.training import moving_averages
from tensorflow.python.training import optimizer
from tensorflow.python.training import saver
+from tensorflow.python.training.saving import saveable_object_util
class MovingAverageOptimizer(optimizer.Optimizer):
@@ -165,7 +166,7 @@
if var_list is None:
var_list = variables.global_variables()
if not isinstance(var_list, dict):
- var_list = saver.BaseSaverBuilder.OpListToDict(var_list)
+ var_list = saveable_object_util.op_list_to_dict(var_list)
v_name_to_tensor = {}
for k, tensor_or_list in six.iteritems(var_list):
diff --git a/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py b/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py
index 200b0d2..8b8065c 100644
--- a/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py
+++ b/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py
@@ -59,6 +59,23 @@
Note that this extension decays weights BEFORE applying the update based
on the gradient, i.e. this extension only has the desired behaviour for
optimizers which do not depend on the value of'var' in the update step!
+
+ Note: when applying a decay to the learning rate, be sure to manually apply
+ the decay to the `weight_decay` as well. For example:
+
+ ```python
+ schedule = tf.train.piecewise_constant(tf.train.get_global_step(),
+ [10000, 15000], [1e-0, 1e-1, 1e-2])
+ lr = 1e-1 * schedule()
+ wd = lambda: 1e-4 * schedule()
+
+ # ...
+
+ optimizer = tf.contrib.opt.MomentumWOptimizer(learning_rate=lr,
+ weight_decay=wd,
+ momentum=0.9,
+ use_nesterov=True)
+ ```
"""
def __init__(self, weight_decay, **kwargs):
diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2.py b/tensorflow/contrib/optimizer_v2/optimizer_v2.py
index 73a556f..7fb23ab 100644
--- a/tensorflow/contrib/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/contrib/optimizer_v2/optimizer_v2.py
@@ -25,6 +25,7 @@
import six
from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribution_strategy_context as distribute_ctx
from tensorflow.python.distribute import reduce_util as ds_reduce_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
@@ -36,7 +37,6 @@
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
-from tensorflow.python.training import distribution_strategy_context as distribute_ctx
from tensorflow.python.training import optimizer as optimizer_v1
from tensorflow.python.training import slot_creator
from tensorflow.python.training.checkpointable import base as checkpointable
@@ -997,10 +997,10 @@
with ops.control_dependencies([update_ops]):
finish_updates = distribution.extended.update_non_slot(
non_slot_devices, finish, group=False)
- # We said grouped=False, which means finish_updates is always a list.
- # It will be [None] when finish() returns None.
- if finish_updates == [None]:
- finish_updates = [update_ops]
+ # We said group=False, which means finish_updates is always a tuple.
+ # It will be (None,) when finish() returns None.
+ if finish_updates == (None,):
+ finish_updates = (update_ops,)
# Update `global_step` (if any).
if global_step is None:
diff --git a/tensorflow/contrib/rate/BUILD b/tensorflow/contrib/rate/BUILD
index c461a71..76db9ae 100644
--- a/tensorflow/contrib/rate/BUILD
+++ b/tensorflow/contrib/rate/BUILD
@@ -34,6 +34,11 @@
name = "rate_test",
size = "small",
srcs = ["rate_test.py"],
+ tags = [
+ "manual", # TODO(b/120555555)
+ "no_oss", # TODO(b/120555555)
+ "notap", # TODO(b/120555555)
+ ],
deps = [
":rate",
"//tensorflow/python:array_ops",
diff --git a/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py b/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py
index d8ca0ea..cec4c3c 100644
--- a/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py
+++ b/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py
@@ -164,6 +164,15 @@
expected = [[[0.0], [27.62]]]
self._assertForwardOpMatchesExpected(input_np, warp_np, expected)
+ expected_grad_data = [[[[0.12], [0.27999997]], [[0.18000001],
+ [0.42000002]]]]
+ expected_grad_warp = [[[0., 0.], [22.60000038, 35.20000076]]]
+
+ grad_output = np.ones([1, 2, 1], dtype=dtype)
+ self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output,
+ expected_grad_data,
+ expected_grad_warp)
+
# One of (x, y) is less than 0.
for dtype in self.float_types:
input_shape = [1, 2, 2, 1]
@@ -171,11 +180,21 @@
input_np = np.array(input_data, dtype=dtype).reshape(input_shape)
warp_shape = [1, 2, 2]
+ # -1 is out of bound for grad_warp.
warp_data = [-1, 0.1, 0.7, 0.6]
warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape)
expected = [[[0.0], [27.62]]]
self._assertForwardOpMatchesExpected(input_np, warp_np, expected)
+ expected_grad_data = [[[[0.12], [0.27999997]], [[0.18000001],
+ [0.42000002]]]]
+ expected_grad_warp = [[[0., 0.], [22.60000038, 35.20000076]]]
+
+ grad_output = np.ones([1, 2, 1], dtype=dtype)
+ self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output,
+ expected_grad_data,
+ expected_grad_warp)
+
# Both of (x, y) are greater than image size.
for dtype in self.float_types:
input_shape = [1, 2, 2, 1]
@@ -183,11 +202,20 @@
input_np = np.array(input_data, dtype=dtype).reshape(input_shape)
warp_shape = [1, 2, 2]
+ # -0.1 is *inbound* for grad_warp and grad_data, 2.1 is out of bound.
warp_data = [-0.1, 0.1, 1.2, 2.1]
warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape)
expected = [[[0.0], [0.0]]]
self._assertForwardOpMatchesExpected(input_np, warp_np, expected)
+ expected_grad_data = [[[[0.81], [0.0]], [[0.09], [0.0]]]]
+ expected_grad_warp = [[[10.30, 2.7], [0.0, 0.0]]]
+
+ grad_output = np.ones([1, 2, 1], dtype=dtype)
+ self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output,
+ expected_grad_data,
+ expected_grad_warp)
+
# One of (x, y) is greater than image size.
for dtype in self.float_types:
input_shape = [1, 2, 2, 1]
@@ -200,6 +228,14 @@
expected = [[[0.0], [0.0]]]
self._assertForwardOpMatchesExpected(input_np, warp_np, expected)
+ expected_grad_data = [[[[0.81], [0.81]], [[0.0], [0.08]]]]
+ expected_grad_warp = [[[-4.5, 9.5], [-9.9, 39.20]]]
+
+ grad_output = np.ones([1, 2, 1], dtype=dtype)
+ self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output,
+ expected_grad_data,
+ expected_grad_warp)
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py
index ffba514..2a4b6ea 100644
--- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py
+++ b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py
@@ -22,53 +22,57 @@
import six
from tensorflow.python.client import session
-from tensorflow.python.estimator import keras as estimator_keras_util
-from tensorflow.python.estimator import model_fn as model_fn_lib
-from tensorflow.python.estimator.export import export as export_helpers
from tensorflow.python.framework import ops
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import models as models_lib
from tensorflow.python.keras import optimizers
from tensorflow.python.keras.engine import sequential
+from tensorflow.python.keras.engine import training_utils
from tensorflow.python.keras.metrics import Metric
from tensorflow.python.keras.models import model_from_json
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import variables
-from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import save as save_lib
from tensorflow.python.saved_model import utils_impl as saved_model_utils
from tensorflow.python.training import saver as saver_lib
from tensorflow.python.training.checkpointable import util as checkpointable_utils
from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow_estimator.python.estimator import keras as estimator_keras_util
+from tensorflow_estimator.python.estimator import model_fn as model_fn_lib
+from tensorflow_estimator.python.estimator.export import export as export_helpers
def save_keras_model(
- model, saved_model_path, custom_objects=None, as_text=None):
- """Save a `tf.keras.Model` into Tensorflow SavedModel format.
+ model, saved_model_path, custom_objects=None, as_text=None,
+ input_signature=None, serving_only=False):
+ """Saves a `tf.keras.Model` into Tensorflow SavedModel format.
`save_model` generates new files/folders under the `saved_model_path` folder:
- 1) an asset folder containing the json string of the model's
- configuration (topology).
- 2) a checkpoint containing the model weights.
- 3) a saved_model.pb file containing the model's MetaGraphs. The prediction
+ 1) a checkpoint containing the model weights.
+ 2) a saved_model.pb file containing the model's MetaGraphs. The prediction
graph is always exported. The evaluaton and training graphs are exported
if the following conditions are met:
- Evaluation: model loss is defined.
- Training: model is compiled with an optimizer defined under `tf.train`.
This is because `tf.keras.optimizers.Optimizer` instances cannot be
saved to checkpoints.
+ 3) Model's json configuration, if model.get_config() has been implemented.
+ This file can be used to reload the model using
+ tf.keras.models.model_from_json(). Note that if any custom objects were
+ used, they should be passed to the `custom_object` argument when loading
+ the model.
- Model Requirements:
- - Model must be a sequential model or functional model. Subclassed models can
- not be saved via this function, unless you provide an implementation for
- get_config() and from_config().
- - All variables must be saveable by the model. In general, this condition is
- met through the use of layers defined in the keras library. However,
- there is currently a bug with variables created in Lambda layer functions
- not being saved correctly (see
- https://github.com/keras-team/keras/issues/9740).
+ Model limitations:
+ - Sequential and functional models can always be saved.
+ - Subclassed models can only be saved when `serving_only=True`. This is due to
+ the current implementation copying the model in order to export the training
+ and evaluation graphs. Because the topology of subclassed models cannot be
+ determined, the subclassed models cannot be cloned. Subclassed models will
+ be entirely exportable in the future.
Note that each mode is exported in separate graphs, so different modes do not
share variables. To use the train graph with evaluation or prediction graphs,
@@ -94,38 +98,88 @@
```
Args:
- model: A `tf.keras.Model` to be saved.
+ model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag
+ `serving_only` must be set to True.
saved_model_path: a string specifying the path to the SavedModel directory.
The SavedModel will be saved to a timestamped folder created within this
directory.
custom_objects: Optional dictionary mapping string names to custom classes
or functions (e.g. custom loss functions).
- as_text: whether to write the `SavedModel` proto in text format.
+ as_text: whether to write the `SavedModel` proto in text format. Currently
+ unavailable in serving-only mode.
+ input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used
+ to specify the expected model inputs. `input_signature`'s nested structure
+ should match the expected nested structure of the inputs to the model. If
+ this is not set, this function will attempt to infer the input shapes and
+ dtypes from the model. Note that if the model is subclassed, the tensor
+ inputs to the call function should be nested in the first argument (this
+ is a general requirement for using subclassed models with Keras functions
+ .fit(), .predict(), etc.).
+ serving_only: Export only the outputs produced from calling the model in
+ predict mode. The losses, optimizer, and other training configurations are
+ not saved. If the SavedModel will only be used for serving (rather than
+ retraining), or if the model is subclassed, this can be set to True.
Returns:
String path to the SavedModel folder, a subdirectory of `saved_model_path`.
Raises:
- NotImplementedError: If the model is a subclassed model.
- ValueError: If a Sequential model does not have input shapes defined by the
- user, and is not built.
+ NotImplementedError: If the model is a subclassed model, and serving_only is
+ False.
+ ValueError: If the input signature cannot be inferred from the model.
"""
+ export_dir = export_helpers.get_timestamped_export_dir(saved_model_path)
+
+ if serving_only:
+ save_lib.save(
+ model, export_dir,
+ signatures=training_utils.trace_model_call(model, input_signature))
+ else:
+ _save_v1_format(model, export_dir, custom_objects, as_text, input_signature)
+
+ try:
+ _export_model_json(model, export_dir)
+ except NotImplementedError:
+ logging.warning('Skipped saving model JSON, subclassed model does not have '
+ 'get_config() defined.')
+
+ return export_dir
+
+
+def _export_model_json(model, saved_model_path):
+ """Saves model configuration as a json string under assets folder."""
+ model_json = model.to_json()
+ model_json_filepath = os.path.join(
+ saved_model_utils.get_or_create_assets_dir(saved_model_path),
+ compat.as_text(constants.SAVED_MODEL_FILENAME_JSON))
+ file_io.write_string_to_file(model_json_filepath, model_json)
+
+
+def _export_model_variables(model, saved_model_path):
+ """Saves model weights in checkpoint format under variables folder."""
+ saved_model_utils.get_or_create_variables_dir(saved_model_path)
+ checkpoint_prefix = saved_model_utils.get_variables_path(saved_model_path)
+ model.save_weights(checkpoint_prefix, save_format='tf', overwrite=True)
+ return checkpoint_prefix
+
+
+def _save_v1_format(model, path, custom_objects, as_text, input_signature):
+ """Exports model to v1 SavedModel format."""
if not model._is_graph_network:
if isinstance(model, sequential.Sequential):
# If input shape is not directly set in the model, the exported model
- # will assume that the inputs have the same shape as the shape the model
- # was built model with.
- if not model.built:
+ # will infer the expected shapes of the input from the model.
+ if not model.built and input_signature is None:
raise ValueError(
- 'Sequential model must be built before it can be exported.')
+ 'Sequential model\'s input shape is unknown. Please build the '
+ 'model, or use the input_signature argument to specify the '
+ 'model inputs.')
else:
raise NotImplementedError(
- 'Exporting subclassed models is not yet supported.')
+ 'Subclassed models can only be exported for serving. Please set '
+ 'argument serving_only=True.')
- export_dir = export_helpers.get_timestamped_export_dir(saved_model_path)
- temp_export_dir = export_helpers.get_temp_export_dir(export_dir)
-
- builder = saved_model_builder._SavedModelBuilder(temp_export_dir)
+ builder = saved_model_builder._SavedModelBuilder(path)
# Manually save variables to export them in an object-based checkpoint. This
# skips the `builder.add_meta_graph_and_variables()` step, which saves a
@@ -133,7 +187,7 @@
# TODO(b/113134168): Add fn to Builder to save with object-based saver.
# TODO(b/113178242): This should only export the model json structure. Only
# one save is needed once the weights can be copied from the model to clone.
- checkpoint_path = _export_model_json_and_variables(model, temp_export_dir)
+ checkpoint_path = _export_model_variables(model, path)
# Export each mode. Use ModeKeys enums defined for `Estimator` to ensure that
# Keras models and `Estimator`s are exported with the same format.
@@ -143,10 +197,12 @@
export_args = {'builder': builder,
'model': model,
'custom_objects': custom_objects,
- 'checkpoint_path': checkpoint_path}
+ 'checkpoint_path': checkpoint_path,
+ 'input_signature': input_signature}
has_saved_vars = False
if model.optimizer:
+ # TODO(kathywu): Verify this works with v2 optimizer.
if isinstance(model.optimizer, optimizers.TFOptimizer):
_export_mode(model_fn_lib.ModeKeys.TRAIN, has_saved_vars, **export_args)
has_saved_vars = True
@@ -161,34 +217,20 @@
builder.save(as_text)
- gfile.Rename(temp_export_dir, export_dir)
- return export_dir
-
-
-def _export_model_json_and_variables(model, saved_model_path):
- """Save model variables and json structure into SavedModel subdirectories."""
- # Save model configuration as a json string under assets folder.
- model_json = model.to_json()
- model_json_filepath = os.path.join(
- saved_model_utils.get_or_create_assets_dir(saved_model_path),
- compat.as_text(constants.SAVED_MODEL_FILENAME_JSON))
- file_io.write_string_to_file(model_json_filepath, model_json)
-
- # Save model weights in checkpoint format under variables folder.
- saved_model_utils.get_or_create_variables_dir(saved_model_path)
- checkpoint_prefix = saved_model_utils.get_variables_path(saved_model_path)
- model.save_weights(checkpoint_prefix, save_format='tf', overwrite=True)
- return checkpoint_prefix
-
def _get_var_list(model):
- """Return list of all checkpointed saveable objects in the model."""
+ """Returns list of all checkpointed saveable objects in the model."""
return checkpointable_utils.named_saveables(model)
+def create_placeholder(spec):
+ return K.placeholder(shape=spec.shape, dtype=spec.dtype, name=spec.name)
+
+
def _export_mode(
- mode, has_saved_vars, builder, model, custom_objects, checkpoint_path):
- """Export a model, and optionally save new vars from the clone model.
+ mode, has_saved_vars, builder, model, custom_objects, checkpoint_path,
+ input_signature):
+ """Exports a model, and optionally saves new vars from the clone model.
Args:
mode: A `tf.estimator.ModeKeys` string.
@@ -199,6 +241,8 @@
custom_objects: A dictionary mapping string names to custom classes
or functions.
checkpoint_path: String path to checkpoint.
+ input_signature: Nested TensorSpec containing the expected inputs. Can be
+ `None`, in which case the signature will be inferred from the model.
Raises:
ValueError: If the train/eval mode is being exported, but the model does
@@ -214,10 +258,16 @@
K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
+ if input_signature is None:
+ input_tensors = None
+ else:
+ input_tensors = nest.map_structure(create_placeholder, input_signature)
+
# Clone the model into blank graph. This will create placeholders for inputs
# and targets.
clone = models_lib.clone_and_build_model(
- model, custom_objects=custom_objects, compile_clone=compile_clone)
+ model, input_tensors=input_tensors, custom_objects=custom_objects,
+ compile_clone=compile_clone)
# Make sure that iterations variable is added to the global step collection,
# to ensure that, when the SavedModel graph is loaded, the iterations
@@ -271,7 +321,7 @@
def _create_signature_def_map(model, mode):
- """Create a SignatureDef map from a Keras model."""
+ """Creates a SignatureDef map from a Keras model."""
inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)}
if model.optimizer:
targets_dict = {x.name.split(':')[0]: x
@@ -309,14 +359,14 @@
def _assert_same_non_optimizer_objects(model, model_graph, clone, clone_graph): # pylint: disable=unused-argument
- """Assert model and clone contain the same checkpointable objects."""
+ """Asserts model and clone contain the same checkpointable objects."""
# TODO(fchollet, kathywu): make sure this works in eager mode.
return True
def load_keras_model(saved_model_path):
- """Load a keras.Model from SavedModel.
+ """Loads a keras.Model from SavedModel.
load_model reinstantiates model state by:
1) loading model topology from json (this will eventually come
diff --git a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py
index 93d73e1..fbf8138 100644
--- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py
+++ b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py
@@ -29,7 +29,9 @@
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.estimator import model_fn as model_fn_lib
+from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.keras.engine import training
from tensorflow.python.keras.utils import tf_utils
@@ -215,7 +217,7 @@
return input_shape
-def functional_model(uses_learning_phase):
+def functional_model(uses_learning_phase=True):
inputs = keras.layers.Input(shape=(3,))
x = keras.layers.Dense(2)(inputs)
x = keras.layers.Dense(3)(x)
@@ -224,7 +226,7 @@
return keras.models.Model(inputs, x)
-def sequential_model(uses_learning_phase):
+def sequential_model(uses_learning_phase=True):
model = keras.models.Sequential()
model.add(keras.layers.Dense(2, input_shape=(3,)))
model.add(keras.layers.Dense(3))
@@ -233,7 +235,7 @@
return model
-def sequential_model_without_input_shape(uses_learning_phase):
+def sequential_model_without_input_shape(uses_learning_phase=True):
model = keras.models.Sequential()
model.add(keras.layers.Dense(2))
model.add(keras.layers.Dense(3))
@@ -242,10 +244,30 @@
return model
+class Subclassed(keras.models.Model):
+
+ def __init__(self):
+ super(Subclassed, self).__init__()
+ self.dense1 = keras.layers.Dense(2)
+ self.dense2 = keras.layers.Dense(3)
+
+ def call(self, inputs):
+ x = self.dense1(inputs)
+ x = self.dense2(x)
+ return x
+
+
+def subclassed_model():
+ return Subclassed()
+
+
def load_model(sess, path, mode):
tags = model_fn_lib.EXPORT_TAG_MAP[mode]
- sig_def_key = (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
- if mode == model_fn_lib.ModeKeys.PREDICT else mode)
+ if mode == model_fn_lib.ModeKeys.PREDICT:
+ sig_def_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+ else:
+ sig_def_key = mode
+
meta_graph_def = loader_impl.load(sess, tags, path)
inputs = {
k: sess.graph.get_tensor_by_name(v.name)
@@ -463,13 +485,54 @@
clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001))
clone.train_on_batch(input_arr, target_arr)
- def testSaveSeqModelWithoutInputShapesRaisesError(self):
- """A Sequential model that hasn't been built should raise an error."""
+ def testSaveSequentialModelWithoutInputShapes(self):
model = sequential_model_without_input_shape(True)
- with self.assertRaisesRegexp(
- ValueError, 'must be built'):
+ # A Sequential model that hasn't been built should raise an error.
+ with self.assertRaisesRegexp(ValueError, 'Please build the model'):
keras_saved_model.save_keras_model(model, '')
+ saved_model_path = self._save_model_dir()
+ output_path = keras_saved_model.save_keras_model(
+ model, saved_model_path,
+ input_signature=tensor_spec.TensorSpec(shape=(10, 11, 12, 13, 14),
+ dtype=dtypes.float32,
+ name='spec_input'))
+
+ with session.Session(graph=ops.Graph()) as sess:
+ inputs, outputs, _ = load_model(sess, output_path,
+ model_fn_lib.ModeKeys.PREDICT)
+ self.assertEqual(5, inputs[next(iter(inputs.keys()))].shape.ndims)
+ self.assertEqual(5, outputs[next(iter(outputs.keys()))].shape.ndims)
+ self.assertEqual(3, outputs[next(iter(outputs.keys()))].shape[-1])
+
+ @test_util.run_v2_only
+ @parameterized.parameters(
+ {
+ 'model_builder': sequential_model_without_input_shape,
+ 'input_signature': [tensor_spec.TensorSpec(shape=[None, 3],
+ dtype=dtypes.float32)]},
+ {
+ 'model_builder': subclassed_model,
+ 'input_signature': [tensor_spec.TensorSpec(shape=[None, 3],
+ dtype=dtypes.float32)]})
+ def testServingOnly(self, model_builder, input_signature):
+ saved_model_path = self._save_model_dir()
+ input_arr = np.random.random((5, 3)).astype(np.float32)
+ model = model_builder()
+ ref_predict = model.predict(input_arr)
+
+ output_path = keras_saved_model.save_keras_model(
+ model, saved_model_path, serving_only=True,
+ input_signature=input_signature)
+
+ # Load predict graph, and test predictions
+ with session.Session(graph=ops.Graph()) as sess:
+ inputs, outputs, _ = load_model(sess, output_path,
+ model_fn_lib.ModeKeys.PREDICT)
+ predictions = sess.run(outputs[next(iter(outputs.keys()))],
+ {inputs[next(iter(inputs.keys()))]: input_arr})
+ self.assertAllClose(ref_predict, predictions, atol=1e-05)
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py b/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py
index 8fcd7ae..f31bdbd 100644
--- a/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py
+++ b/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py
@@ -19,6 +19,7 @@
import numpy as np
+from tensorflow.python import tf2
from tensorflow.contrib.solvers.python.ops import lanczos
from tensorflow.contrib.solvers.python.ops import util
from tensorflow.python.framework import constant_op
@@ -80,7 +81,8 @@
for shape in [[4, 4], [7, 4], [5, 8]]:
for orthogonalize in True, False:
for steps in range(1, min(shape) + 1):
- for use_static_shape in True, False:
+ # TF2 does not support placeholders so we skip it
+ for use_static_shape in set([True, tf2.enabled()]):
arg_string = "%s_%s_%s_%s_staticshape_%s" % (
dtype.__name__, "_".join(map(str, shape)), orthogonalize, steps,
use_static_shape)
diff --git a/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py b/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py
index 2a91009..841a41a 100644
--- a/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py
+++ b/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py
@@ -19,6 +19,7 @@
import numpy as np
+from tensorflow.python import tf2
from tensorflow.contrib.solvers.python.ops import least_squares
from tensorflow.contrib.solvers.python.ops import util
from tensorflow.python.framework import constant_op
@@ -76,7 +77,8 @@
if __name__ == "__main__":
for dtype in np.float32, np.float64:
for shape in [[4, 4], [8, 5], [3, 7]]:
- for use_static_shape in True, False:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_static_shape in set([True, tf2.enabled()]):
arg_string = "%s_%s_staticshape_%s" % (dtype.__name__,
"_".join(map(str, shape)),
use_static_shape)
diff --git a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py
index a0e6eb8..10807f7 100644
--- a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py
+++ b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py
@@ -19,6 +19,7 @@
import numpy as np
+from tensorflow.python import tf2
from tensorflow.contrib.solvers.python.ops import linear_equations
from tensorflow.contrib.solvers.python.ops import util
from tensorflow.python.framework import constant_op
@@ -113,7 +114,8 @@
if __name__ == "__main__":
for dtype in np.float32, np.float64:
for size in 1, 4, 10:
- for use_static_shape in True, False:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_static_shape in set([True, tf2.enabled()]):
shape = [size, size]
arg_string = "%s_%s_staticshape_%s" % (dtype.__name__, size,
use_static_shape)
diff --git a/tensorflow/contrib/tensorrt/README.md b/tensorflow/contrib/tensorrt/README.md
index caf8b6d..a9c2ad7 100644
--- a/tensorflow/contrib/tensorrt/README.md
+++ b/tensorflow/contrib/tensorrt/README.md
@@ -1,8 +1,46 @@
-# Using TensorRT in TensorFlow
+# Using TensorRT in TensorFlow (TF-TRT)
-This module provides necessary bindings and introduces TRT_engine_op operator
-that wraps a subgraph in TensorRT. This is still a work in progress but should
-be useable with most common graphs.
+This module provides necessary bindings and introduces `TRTEngineOp` operator
+that wraps a subgraph in TensorRT. This module is under active development.
+
+## Installing TF-TRT
+
+Currently TensorFlow nightly builds include TF-TRT by default, which means you
+don't need to install TF-TRT separately. You can pull the latest TF containers
+from docker hub or install the latest TF pip package to get access to the latest
+TF-TRT.
+
+If you want to use TF-TRT on NVIDIA Jetson platform, you can find the download
+links for the relevant TensorFlow pip packages here:
+https://docs.nvidia.com/deeplearning/dgx/index.html#installing-frameworks-for-jetson
+
+## Installing TensorRT
+
+In order to make use of TF-TRT, you will need a local installation of TensorRT.
+Installation instructions for compatibility with TensorFlow are provided on the
+[TensorFlow GPU support](https://www.tensorflow.org/install/gpu) guide.
+
+## Examples
+
+You can find example scripts for running inference on deep learning models in
+this repository: https://github.com/tensorflow/tensorrt
+
+We have used these examples to verify the accuracy and performance of TF-TRT.
+For more information see
+[Verified Models](https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html#verified-models).
+
+## Documentation
+
+[TF-TRT documentation](https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html)
+gives an overview of the supported functionalities, provides tutorials and
+verified models, explains best practices with troubleshooting guides.
+
+## Tests
+
+TF-TRT includes both Python tests and C++ unit tests. Most of Python tests are
+located in the test directory and they can be executed using `bazel test` or
+directly with the Python command. Most of the C++ unit tests are used to test
+the conversion functions that convert each TF op to a number of TensorRT layers.
## Compilation
@@ -18,12 +56,3 @@
bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/
```
-After the installation of tensorflow package, TensorRT transformation will be
-available. An example use can be found in test/test_tftrt.py script
-
-## Installing TensorRT 3.0.4
-
-In order to make use of TensorRT integration, you will need a local installation
-of TensorRT 3.0.4 from the [NVIDIA Developer website](https://developer.nvidia.com/tensorrt).
-Installation instructions for compatibility with TensorFlow are provided on the
-[TensorFlow GPU support](https://www.tensorflow.org/install/gpu) guide.
diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index 3b32f72..623cd79 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -89,49 +89,52 @@
// TODO(laigd): move this set to TrtNodeValidator where it should belong.
// LINT.IfChange
static const std::set<string> candidate_ops = {
- "Identity",
- "Snapshot",
+ "Abs",
+ "Add",
+ "AvgPool",
+ "BatchMatMul",
+ "BiasAdd",
+ "ConcatV2",
"Const",
"Conv2D",
- "MaxPool",
- "BiasAdd",
- "Relu",
- "Sigmoid",
- "Tanh",
- "Add",
- "Mul",
- "Sub",
- "Rsqrt",
- "Pad",
- "Mean",
- "AvgPool",
- "ConcatV2",
"DepthwiseConv2dNative",
+ "Div",
+ "Exp",
+ "ExpandDims",
"FusedBatchNorm",
"FusedBatchNormV2",
- "Div",
- "RealDiv",
- "Rsqrt",
- "Reciprocal",
- "Exp",
+ "Identity",
"Log",
- "Sqrt",
- "Abs",
- "Neg",
- "Transpose",
- "Reshape",
"MatMul",
- "BatchMatMul",
- "Softmax",
- "Minimum",
- "Maximum",
- "TopKV2",
- "Sum",
- "Prod",
"Max",
+ "MaxPool",
+ "Maximum",
+ "Mean",
"Min",
+ "Minimum",
+ "Mul",
+ "Neg",
+ "Pad",
+ "Prod",
+ "RealDiv",
+ "Reciprocal",
+ "Relu",
"Relu6",
+ "Reshape",
+ "Rsqrt",
+ "Rsqrt",
+ "Sigmoid",
+ "Snapshot",
+ "Softmax",
+ "Sqrt",
"Square",
+ "Squeeze",
+ "StridedSlice",
+ "Sub",
+ "Sum",
+ "Tanh",
+ "TopKV2",
+ "Transpose",
};
bool is_supported_op_type =
(candidate_ops.count(node->type_string()) ||
@@ -585,6 +588,14 @@
}
}
}
+ // We don't support segments with no inputs. Fall back to native TF here to
+ // avoid crash later. Constant folding should've folded the ops that make up
+ // these segments.
+ if (inputs.empty()) {
+ return tensorflow::errors::Internal(
+ "Segment has no inputs (possible "
+ "constfold failure)");
+ }
const bool calibrate_int8 =
(info.precision_mode == INT8MODE && info.use_calibration);
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index fee0956..adf8831 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -120,6 +120,15 @@
return trt_dims;
}
+Status TensorShapeArrayToTrtDims(const std::vector<int>& shape,
+ nvinfer1::Dims* out,
+ bool ignore_first_dim = false) {
+ PartialTensorShape tensor_shape;
+ TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(shape, &tensor_shape));
+ *out = TensorShapeToTrtDims(tensor_shape, ignore_first_dim);
+ return tensorflow::Status::OK();
+}
+
void GetOutputProperties(const grappler::GraphProperties& graph_properties,
const Node* node, const int out_port,
PartialTensorShape* shape,
@@ -623,6 +632,11 @@
return this->at(key)->b();
}
+template <>
+int TFAttrs::get<int>(const string& key) const {
+ return this->at(key)->i();
+}
+
// TODO(jie): reorder4 & reorder2 should be merged?
// TODO(aaroey): fix the order of parameters.
template <typename T>
@@ -1524,6 +1538,24 @@
tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) {
const auto& inputs = params->inputs;
const auto& node_def = params->node_def;
+ if (inputs.at(0).is_weights()) {
+ return tensorflow::errors::Unimplemented(
+ node_def.op(), " is only implemented for tensors, not weights, at ",
+ node_def.name());
+ }
+ if (inputs.at(1).is_tensor()) {
+ return tensorflow::errors::Unimplemented("Kernel for ", node_def.op(),
+ " must be constant weights, at ",
+ node_def.name());
+ }
+ TRT_ShapedWeights weights_rsck = inputs.at(1).weights();
+ VLOG(2) << "weight shape: " << weights_rsck.DebugString();
+ if (weights_rsck.shape_.nbDims != 4) {
+ return tensorflow::errors::Internal(
+ "Conv2D expects kernel of dimension 4, at: " + node_def.name());
+ }
+ if (params->validation_only) return tensorflow::Status::OK();
+
const nvinfer1::ITensor* tensor = inputs.at(0).tensor();
TFAttrs attrs(node_def);
@@ -1545,12 +1577,6 @@
if (num_groups == 0) num_groups = tensor_dim.d[0]; // depthwise convolution
VLOG(2) << "groups count: " << num_groups;
- TRT_ShapedWeights weights_rsck = inputs.at(1).weights();
- VLOG(2) << "weight shape: " << weights_rsck.DebugString();
- if (weights_rsck.shape_.nbDims != 4) {
- return tensorflow::errors::Internal(
- "Conv2D expects kernel of dimension 4, at: " + node_def.name());
- }
if (params->converter->precision_mode() == FP16MODE) {
weights_rsck =
ConvertFP32ToFP16(params->weight_store, inputs.at(1).weights());
@@ -1637,7 +1663,7 @@
case ConvolutionType::DEPTHWISE_CONV:
return ConvertConv2DHelper(params, 0);
}
- return tensorflow::errors::Unimplemented("unsupported convolution type at, " +
+ return tensorflow::errors::Unimplemented("Unsupported convolution type, at ",
params->node_def.name());
}
@@ -1880,6 +1906,372 @@
return tensorflow::Status::OK();
}
+tensorflow::Status ConvertExpandDims(OpConverterParams* params) {
+ const auto& inputs = params->inputs;
+ const auto& node_def = params->node_def;
+ if (inputs.size() != 2) {
+ return tensorflow::errors::InvalidArgument(
+ "Two inputs expected for ExpandDims, at ", node_def.name());
+ }
+ if (inputs.at(0).is_weights()) {
+ return tensorflow::errors::Unimplemented(
+ "ExpandDims expects tensor for input, at ", node_def.name());
+ }
+ if (!inputs.at(1).is_weights()) {
+ return tensorflow::errors::InvalidArgument(
+ "ExpandDims expects weights for axis, at ", node_def.name());
+ }
+ // Get input shape as vector.
+ TRT_TensorOrWeights input_tensor = inputs.at(0);
+ const nvinfer1::Dims dims = input_tensor.GetTrtDims();
+ std::vector<int> input_dims(dims.d, dims.d + dims.nbDims);
+ // Add batch dim back.
+ input_dims.insert(input_dims.begin(), -1);
+ const int input_rank = input_dims.size();
+ // Get axis to expand on.
+ TRT_ShapedWeights weights = inputs.at(1).weights();
+ if (weights.count() != 1) {
+ return tensorflow::errors::InvalidArgument(
+ "ExpandDims axis must be a scalar, at ", node_def.name());
+ }
+ const int* weights_ptr =
+ static_cast<int*>(const_cast<void*>(weights.GetValues()));
+ int axis = weights_ptr[0];
+ // Make sure axis is valid.
+ if ((axis < (-input_rank - 1)) || (axis > input_rank)) {
+ return tensorflow::errors::InvalidArgument(
+ "Axis for ExpandDims is invalid, must be in the range "
+ "[-rank(input) - 1, rank(input)], at ",
+ node_def.name());
+ }
+ // Convert negative axis to corresponding positive axis.
+ if (axis < 0) axis += input_rank + 1;
+ if (axis == 0) {
+ return tensorflow::errors::Unimplemented(
+ "Modifying batch dimension is not supported for ExpandDims, at ",
+ node_def.name());
+ }
+ if (params->validation_only) return Status::OK();
+
+ // ExpandDims: Insert new dim of size 1.
+ input_dims.insert(input_dims.begin() + axis, 1);
+ // Reshape tensor.
+ nvinfer1::Dims new_dims;
+ TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &new_dims,
+ /*ignore_first_dim=*/true));
+ const nvinfer1::ITensor* output_tensor = nullptr;
+ TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape(
+ input_tensor, new_dims, &output_tensor));
+ params->outputs->push_back(
+ TRT_TensorOrWeights(const_cast<nvinfer1::ITensor*>(output_tensor)));
+ return tensorflow::Status::OK();
+}
+
+tensorflow::Status ConvertSqueeze(OpConverterParams* params) {
+ const auto& inputs = params->inputs;
+ const auto& node_def = params->node_def;
+ if (inputs.size() != 1) {
+ return tensorflow::errors::InvalidArgument(
+ "One input expected for Squeeze, at ", node_def.name());
+ }
+ if (inputs.at(0).is_weights()) {
+ return tensorflow::errors::Unimplemented(
+ "Squeeze expects tensor for input, at ", node_def.name());
+ }
+ // Get input shape.
+ TRT_TensorOrWeights input_tensor = inputs.at(0);
+ const nvinfer1::Dims dims = input_tensor.GetTrtDims();
+ std::vector<int> input_dims(dims.d, dims.d + dims.nbDims);
+ // Add batch dim back.
+ input_dims.insert(input_dims.begin(), -1);
+ const int input_rank = input_dims.size();
+ // Mark axes to remove by setting them to 0.
+ TFAttrs attrs(node_def);
+ auto squeeze_dims = attrs.get<std::vector<int>>("squeeze_dims");
+ if (squeeze_dims.size() == 0) {
+ return tensorflow::errors::Unimplemented(
+ "Squeeze is only implemented for explicit dims, at ", node_def.name());
+ }
+ for (int axis : squeeze_dims) {
+ // Make sure axis is valid.
+ if ((axis < -input_rank) || (axis >= input_rank)) {
+ return tensorflow::errors::InvalidArgument(
+ "Axis for Squeeze is invalid, must be in the range "
+ "[-rank(input), rank(input)), at ",
+ node_def.name());
+ }
+ // Convert negative axis to corresponding positive axis.
+ if (axis < 0) axis += input_rank;
+ // Don't squeeze batch dim.
+ if (axis == 0) {
+ return tensorflow::errors::Unimplemented(
+ "Cannot squeeze batch dimension, at ", node_def.name());
+ }
+ // Make sure target dimension is size 1.
+ if (input_dims[axis] != 1) {
+ return tensorflow::errors::InvalidArgument(
+ "Cannot squeeze a dimension which isn't size 1, at ",
+ node_def.name());
+ }
+ // Mark dim for removal by setting to 0.
+ input_dims[axis] = 0;
+ }
+ if (params->validation_only) return Status::OK();
+
+ // Remove all dims which are equal to 0.
+ input_dims.erase(std::remove(input_dims.begin(), input_dims.end(), 0),
+ input_dims.end());
+ // Reshape tensor.
+ nvinfer1::Dims new_dims;
+ TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &new_dims,
+ /*ignore_first_dim=*/true));
+ const nvinfer1::ITensor* output_tensor = nullptr;
+ TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape(
+ input_tensor, new_dims, &output_tensor));
+ params->outputs->push_back(
+ TRT_TensorOrWeights(const_cast<nvinfer1::ITensor*>(output_tensor)));
+ return tensorflow::Status::OK();
+}
+
+// Gets the bounds (start or end) from the weights of a StridedSlice op.
+tensorflow::Status GetStridedSliceBound(const std::vector<int>& input_dims,
+ const TRT_ShapedWeights& bound_weights,
+ int mask, bool begin, string node_name,
+ std::vector<int>* output_bound) {
+ const string bound_name = (begin) ? "begin" : "end";
+ const int* weights_ptr = static_cast<int*>(bound_weights.GetValues());
+ *output_bound =
+ std::vector<int>(weights_ptr, weights_ptr + bound_weights.count());
+ if (output_bound->size() != input_dims.size()) {
+ return tensorflow::errors::InvalidArgument(
+ "StridedSlice \"", bound_name, "\" specified ",
+ std::to_string(output_bound->size()), " dimensions, but input rank is ",
+ std::to_string(input_dims.size()), ", at ", node_name);
+ }
+ for (int i = 0; i < output_bound->size(); i++) {
+ if ((1 << i) & mask) {
+ // Apply mask.
+ (*output_bound)[i] = (begin) ? 0 : input_dims[i];
+ // Masked bound will always result in a valid, non-negative bound, so we
+ // don't need the following checks. For the common case of using masks on
+ // a undefined batch dim (-1), we specifically don't want to do the
+ // following checks because they will erroneously detect an out of range
+ // bound or try to correct the negative value.
+ continue;
+ }
+ // Make sure bound is valid.
+ if (((*output_bound)[i] < -input_dims[i]) ||
+ ((*output_bound)[i] > input_dims[i])) {
+ return tensorflow::errors::InvalidArgument(
+ bound_name, " value of ", std::to_string((*output_bound)[i]),
+ " for StridedSlice is invalid, must be in the range "
+ "[-dim_size(i), dim_size(i)], at ",
+ node_name);
+ }
+ // Convert negative values to their positive equivalent.
+ if ((*output_bound)[i] < 0) {
+ (*output_bound)[i] += input_dims[i];
+ }
+ }
+ return tensorflow::Status::OK();
+}
+
+tensorflow::Status ConvertStridedSlice(OpConverterParams* params) {
+ const auto& inputs = params->inputs;
+ const auto& node_def = params->node_def;
+ if (inputs.size() != 4) {
+ return tensorflow::errors::InvalidArgument(
+ "StridedSlice expects 4 inputs, at ", node_def.name());
+ }
+ if (!inputs.at(1).is_weights() || !inputs.at(2).is_weights() ||
+ !inputs.at(3).is_weights()) {
+ return tensorflow::errors::InvalidArgument(
+ "StridedSlice expects weights for begin, end, and strides, at ",
+ node_def.name());
+ }
+ if (!inputs.at(0).is_tensor()) {
+ return tensorflow::errors::Unimplemented(
+ "StridedSlice is only implemented for tensors, at ", node_def.name());
+ }
+ // Get input dims.
+ nvinfer1::Dims dims = inputs.at(0).GetTrtDims();
+ std::vector<int> input_dims(dims.d, dims.d + dims.nbDims);
+ if (inputs.at(0).is_tensor()) {
+ // Temporarily add batch dimension so that indexes line up properly.
+ input_dims.insert(input_dims.begin(), inputs.at(0).batch_size());
+ }
+ if (input_dims.size() > 4) {
+ return tensorflow::errors::Unimplemented(
+ "StridedSlice is not implemented for tensors with rank > 4, at ",
+ node_def.name());
+ }
+ TFAttrs attrs(node_def);
+ // Get begin and end bounds per axis.
+ std::vector<int> begin, end;
+ TF_RETURN_IF_ERROR(GetStridedSliceBound(input_dims, inputs.at(1).weights(),
+ attrs.get<int>("begin_mask"), true,
+ node_def.name(), &begin));
+ TF_RETURN_IF_ERROR(GetStridedSliceBound(input_dims, inputs.at(2).weights(),
+ attrs.get<int>("end_mask"), false,
+ node_def.name(), &end));
+ // Get strides per axis (must all be 1).
+ TRT_ShapedWeights stride_weights = inputs.at(3).weights();
+ const int* stride_weights_ptr = static_cast<int*>(stride_weights.GetValues());
+ std::vector<int> strides(stride_weights_ptr,
+ stride_weights_ptr + stride_weights.count());
+ for (int x : strides) {
+ if (x != 1) {
+ return tensorflow::errors::Unimplemented(
+ "StridedSlice is only implemented for stride of 1, at ",
+ node_def.name());
+ }
+ }
+ // Unsupported mask options.
+ for (const string& attr :
+ {"ellipsis_mask", "new_axis_mask", "shrink_axis_mask"}) {
+ int attr_val = attrs.get<int>(attr);
+ if (attr_val != 0) {
+ return tensorflow::errors::Unimplemented(
+ attr, " is not supported for StridedSlice, at ", node_def.name());
+ }
+ }
+
+ nvinfer1::ITensor* tensor =
+ const_cast<nvinfer1::ITensor*>(inputs.at(0).tensor());
+ // Reshape if necessary to 4-D, since IPaddingLayer requires a 4-D input.
+ const bool need_reshape = (input_dims.size() != 4);
+ int reshape_dims_added = 0;
+ nvinfer1::Dims reshape_dims;
+ if (need_reshape) {
+ // Add new dims after batch dim until tensor is 4D.
+ while (input_dims.size() < 4) {
+ input_dims.insert(input_dims.begin() + 1, 1);
+ begin.insert(begin.begin() + 1, 0);
+ end.insert(end.begin() + 1, 1);
+ reshape_dims_added++;
+ }
+ TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &reshape_dims,
+ /*ignore_first_dim=*/true));
+ }
+ // Find dimensions which need to be sliced.
+ std::vector<int> pad_dims;
+ for (int i = 0; i < input_dims.size(); i++) {
+ if ((begin[i] != 0) || (end[i] != input_dims[i])) {
+ if (i == 0) {
+ return tensorflow::errors::Unimplemented(
+ "StridedSlice can't modify batch dim, at ", node_def.name());
+ } else if ((end[i] - begin[i]) < 0) {
+ return tensorflow::errors::InvalidArgument(
+ "New size of sliced dimension is negative, at ", node_def.name());
+ }
+ pad_dims.push_back(i);
+ }
+ }
+ if (pad_dims.size() == 0) {
+ // No dimensions are changed. We could create a padding layer anyway with
+ // values of 0.
+ if (params->validation_only) return Status::OK();
+ params->outputs->push_back(inputs.at(0));
+ return tensorflow::Status::OK();
+ } else if (pad_dims.size() == 1) {
+ // Only one dim is modified but we have to have 2, mark a second dim which
+ // will have padding of 0. The dim we add is chosen to avoid an unecessary
+ // transpose.
+ if (pad_dims[0] != 2) {
+ pad_dims.push_back(2);
+ } else {
+ pad_dims.push_back(3);
+ }
+ } else if (pad_dims.size() > 2) {
+ return tensorflow::errors::Unimplemented(
+ "StridedSlice can only modify 2 dimensions, at ", node_def.name());
+ }
+ std::sort(pad_dims.begin(), pad_dims.end());
+ // Convert to pre/post padding values. Since TRT does not have a StridedSlice
+ // or Slice layer, we instead create an IPaddingLayer with negative padding.
+ nvinfer1::DimsHW pre_padding, post_padding;
+ for (int i = 0; i < pad_dims.size(); i++) {
+ const int axis = pad_dims[i];
+ pre_padding.d[i] = -begin[axis];
+ post_padding.d[i] = end[axis] - input_dims[axis];
+ }
+
+ // IPaddingLayer will always apply the padding to dims 2,3 (input format is
+ // NCHW).
+ const bool need_transpose = !(pad_dims[0] == 2 && pad_dims[1] == 3);
+ std::vector<int> transpose_order(input_dims.size());
+ std::vector<int> inv_transpose_order(input_dims.size());
+ if (need_transpose) {
+ if (pad_dims[0] == 1 && pad_dims[1] == 3) {
+ transpose_order = {0, 2, 1, 3};
+ inv_transpose_order = {0, 2, 1, 3};
+ } else if (pad_dims[0] == 1 && pad_dims[1] == 2) {
+ transpose_order = {0, 3, 1, 2};
+ inv_transpose_order = {0, 2, 3, 1};
+ }
+ }
+ if (params->validation_only) return Status::OK();
+
+ // Start conversion.
+ if (need_reshape) {
+ const nvinfer1::ITensor* output_tensor = nullptr;
+ TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape(
+ inputs.at(0), reshape_dims, &output_tensor));
+ tensor = const_cast<nvinfer1::ITensor*>(output_tensor);
+ }
+ if (need_transpose) {
+ const nvinfer1::ITensor* output_tensor = nullptr;
+ TF_RETURN_IF_ERROR(params->converter->TransposeTensor(
+ tensor, transpose_order, &output_tensor));
+ tensor = const_cast<nvinfer1::ITensor*>(output_tensor);
+ }
+
+ // Add padding layer
+ nvinfer1::IPaddingLayer* layer = params->converter->network()->addPadding(
+ *const_cast<nvinfer1::ITensor*>(tensor), pre_padding, post_padding);
+ TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
+ params->converter->MarkQuantizationRangesAsInferrable(tensor,
+ layer->getOutput(0));
+ tensor = layer->getOutput(0);
+
+ // Restore transpose
+ if (need_transpose) {
+ const nvinfer1::ITensor* output_tensor = nullptr;
+ TF_RETURN_IF_ERROR(params->converter->TransposeTensor(
+ tensor, inv_transpose_order, &output_tensor));
+ tensor = const_cast<nvinfer1::ITensor*>(output_tensor);
+ }
+ // Restore reshape
+ if (need_reshape) {
+ // Calculate output dimensions
+ for (int i = 0; i < pad_dims.size(); i++) {
+ const int axis = pad_dims[i];
+ input_dims[axis] = end[axis] - begin[axis];
+ }
+ // Remove added 1 dimensions
+ for (int i = 0; i < reshape_dims_added; i++) {
+ int value = input_dims[1];
+ if (value != 1) {
+ return tensorflow::errors::Internal(
+ "StridedSlice error when reshaping, at ", node_def.name());
+ }
+ input_dims.erase(input_dims.begin() + 1);
+ }
+
+ nvinfer1::Dims new_dims;
+ TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &new_dims,
+ /*ignore_first_dim=*/true));
+ const nvinfer1::ITensor* output_tensor = nullptr;
+ TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape(
+ TRT_TensorOrWeights(tensor), new_dims, &output_tensor));
+ tensor = const_cast<nvinfer1::ITensor*>(output_tensor);
+ }
+
+ params->outputs->push_back(
+ TRT_TensorOrWeights(const_cast<nvinfer1::ITensor*>(tensor)));
+ return tensorflow::Status::OK();
+}
+
tensorflow::Status ConvertConv2D(OpConverterParams* params) {
return ConvertConv2DHelper(params, ConvolutionType::DEFAULT);
}
@@ -1891,9 +2283,29 @@
tensorflow::Status ConvertPool(OpConverterParams* params) {
const auto& inputs = params->inputs;
const auto& node_def = params->node_def;
- const nvinfer1::ITensor* tensor = inputs.at(0).tensor();
+ if (inputs.at(0).is_weights()) {
+ return tensorflow::errors::Unimplemented(
+ node_def.op(), " is only implemented for tensors, not weights, at ",
+ node_def.name());
+ }
+ nvinfer1::PoolingType type;
+ if (node_def.op() == "MaxPool") {
+ type = nvinfer1::PoolingType::kMAX;
+ } else if (node_def.op() == "AvgPool") {
+ type = nvinfer1::PoolingType::kAVERAGE;
+ } else {
+ return tensorflow::errors::Unimplemented(
+ "Unsupported pooling type: ", node_def.op(), ", at ", node_def.name());
+ }
TFAttrs attrs(node_def);
+ const string padding_type = attrs.get<string>("padding");
+ if ((padding_type != "SAME") && (padding_type != "VALID")) {
+ return tensorflow::errors::Unimplemented(
+ "Unsupported padding type: ", padding_type, ", at ", node_def.name());
+ }
+ if (params->validation_only) return Status::OK();
+ const nvinfer1::ITensor* tensor = inputs.at(0).tensor();
int h_index = 2;
int w_index = 3;
const auto data_format = attrs.get<string>("data_format");
@@ -1904,16 +2316,6 @@
const_cast<nvinfer1::ITensor*>(tensor), {0, 3, 1, 2}, &tensor));
}
- nvinfer1::PoolingType type;
- if (node_def.op() == "MaxPool") {
- type = nvinfer1::PoolingType::kMAX;
- } else if (node_def.op() == "AvgPool") {
- type = nvinfer1::PoolingType::kAVERAGE;
- } else {
- return tensorflow::errors::Unimplemented("Unsupported pool type: ",
- node_def.op());
- }
-
const auto tf_stride = attrs.get<std::vector<int>>("strides");
const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]);
@@ -1922,7 +2324,6 @@
auto tensor_dim = tensor->getDimensions();
std::vector<std::pair<int, int>> padding;
- const string padding_type = attrs.get<string>("padding");
if (padding_type == "SAME") {
// This is NCHW tensor with no batch dimension.
// 1 -> h
@@ -1932,9 +2333,6 @@
{static_cast<int>(tensor_dim.d[1]), static_cast<int>(tensor_dim.d[2])});
} else if (padding_type == "VALID") {
padding = {{0, 0}, {0, 0}};
- } else {
- return tensorflow::errors::Unimplemented("Unsupported padding type: ",
- padding_type);
}
if (padding[0].first != padding[0].second ||
@@ -2701,6 +3099,7 @@
return tensorflow::errors::Unimplemented(
"Padding layer does not support padding on dimension 1 and 3 yet");
}
+ if (params->validation_only) return Status::OK();
bool legit_pad = true;
nvinfer1::DimsHW pre_padding(0, 0);
@@ -2804,6 +3203,7 @@
inputs_vec.push_back(tensor_i);
}
+ if (params->validation_only) return tensorflow::Status::OK();
// nvinfer1::ITensor const* tensor = inputs.at(0).tensor();
nvinfer1::IConcatenationLayer* layer =
@@ -2825,12 +3225,35 @@
auto data_format = attrs.get<string>("data_format");
if (data_format != "NCHW") {
return tensorflow::errors::Unimplemented(
- "only data_format=NCHW is supported, at " + node_def.name());
+ node_def.op(), " only supports data_format=NCHW, at ", node_def.name());
}
bool is_training = attrs.get<bool>("is_training");
if (is_training) {
+ // 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();
return tensorflow::errors::Unimplemented(
- "only is_training=false is supported, at " + node_def.name());
+ node_def.op(), " only supports is_training=false, at ",
+ node_def.name());
+ }
+ if (inputs.at(0).is_weights()) {
+ return tensorflow::errors::Unimplemented(
+ node_def.op(),
+ " is only implemented for tensor inputs, not weights, at ",
+ node_def.name());
+ }
+ for (int i = 1; i < 5; i++) {
+ if (inputs.at(i).is_tensor()) {
+ return tensorflow::errors::Unimplemented(
+ node_def.op(),
+ " must have constant inputs for scale, offset, mean and variance, "
+ "at ",
+ node_def.name());
+ }
}
nvinfer1::ITensor const* tensor = inputs.at(0).tensor();
@@ -2845,7 +3268,7 @@
for (int i = 1; i < 5; i++) {
if (inputs.at(i).weights().type_ != parameter_type) {
return tensorflow::errors::Unimplemented(
- "Inconsistent parameter type for batchnormis not supported, at: " +
+ "Inconsistent parameter type for batchnorm is not supported, at: " +
node_def.name());
}
}
@@ -2865,6 +3288,8 @@
"Inconsistent batchnorm parameter count, at: " + node_def.name());
}
}
+ if (params->validation_only) return Status::OK();
+
// We could technically have two weights with different shape.
// that requires two addScale op, arguably less performant
TRT_ShapedWeights combined_scale_weights =
@@ -3150,12 +3575,19 @@
std::unordered_map<string, OpConverter>* registration) {
// TODO(laigd): support all op types.
(*registration)["BiasAdd"] = ConvertBiasAdd;
+ (*registration)["ConcatV2"] = ConvertConcat;
(*registration)["Const"] = ConvertConst;
- (*registration)["Transpose"] = ConvertTranspose;
- (*registration)["Reshape"] = ConvertReshape;
+ (*registration)["Conv2D"] = ConvertConv2D;
+ (*registration)["DepthwiseConv2dNative"] = ConvertConv2DDepthwise;
+ (*registration)["ExpandDims"] = ConvertExpandDims;
(*registration)["MatMul"] = ConvertMatMul;
+ (*registration)["Pad"] = ConvertPad;
(*registration)["Relu6"] = ConvertRelu6;
+ (*registration)["Reshape"] = ConvertReshape;
(*registration)["Square"] = ConvertSquare;
+ (*registration)["Squeeze"] = ConvertSqueeze;
+ (*registration)["StridedSlice"] = ConvertStridedSlice;
+ (*registration)["Transpose"] = ConvertTranspose;
for (auto quantization_op_type :
{"QuantizeAndDequantizeV2", "QuantizeAndDequantizeV3",
@@ -3169,6 +3601,12 @@
for (auto activation_op_type : {"Relu", "Sigmoid", "Tanh"}) {
(*registration)[activation_op_type] = ConvertActivation;
}
+ for (auto pool_op_type : {"AvgPool", "MaxPool"}) {
+ (*registration)[pool_op_type] = ConvertPool;
+ }
+ for (auto normalization_op_type : {"FusedBatchNorm", "FusedBatchNormV2"}) {
+ (*registration)[normalization_op_type] = ConvertFusedBatchNorm;
+ }
}
void TrtNodeValidator::RegisterOpValidators() {
@@ -3177,21 +3615,10 @@
void Converter::RegisterOpConverters() {
RegisterValidatableOpConverters(&op_registry_);
-
- op_registry_["Conv2D"] = ConvertConv2D;
- op_registry_["DepthwiseConv2dNative"] = ConvertConv2DDepthwise;
- op_registry_["MaxPool"] = ConvertPool;
- op_registry_["AvgPool"] = ConvertPool;
// TODO(ben,jie): this is a temp hack.
op_registry_["Identity"] = ConvertIdentity; // Identity should be removed
op_registry_["Snapshot"] = ConvertIdentity; // Snapshot should be removed
- op_registry_["Pad"] = ConvertPad;
-
- op_registry_["ConcatV2"] = ConvertConcat;
- op_registry_["FusedBatchNorm"] = ConvertFusedBatchNorm;
- op_registry_["FusedBatchNormV2"] = ConvertFusedBatchNorm;
-
op_registry_["Rsqrt"] = ConvertUnary;
op_registry_["Reciprocal"] = ConvertUnary;
op_registry_["Exp"] = ConvertUnary;
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc
index 4430333..a2ddfbf 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc
@@ -2113,6 +2113,512 @@
}
}
+TEST_F(OpConverterTest, ConvertExpandDims) {
+ {
+ // Input list is empty, should fail.
+ NodeDef node_def = MakeNodeDef("my_expanddims", "ExpandDims", {});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "Two inputs expected for ExpandDims, at my_expanddims");
+ }
+
+ // Get the NodeDef for ExpandDims.
+ Scope s = Scope::NewRootScope();
+ auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT);
+ auto weights = ops::Placeholder(s.WithOpName("weights"), DT_INT32);
+ auto expanddims =
+ ops::ExpandDims(s.WithOpName("my_expanddims"), input, weights);
+ const NodeDef& node_def = expanddims.operation.node()->def();
+ {
+ // Input is weights, should fail.
+ Reset();
+ AddTestWeights<int32>("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6});
+ AddTestWeights<int32>("weights", {1}, {1});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "ExpandDims expects tensor for input, at my_expanddims");
+ }
+ {
+ // Axis is a tensor, should fail.
+ Reset();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestTensor("weights", {3});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "ExpandDims expects weights for axis, at my_expanddims");
+ }
+ {
+ // Add dim at batch dimension, should fail.
+ Reset();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("weights", {1}, {0});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "Modifying batch dimension is not supported for ExpandDims, at "
+ "my_expanddims");
+ }
+ {
+ // Add dim at batch dimension via negative axis, should fail.
+ Reset();
+ AddTestTensor("input", {1, 2, 3});
+ // Input is rank 4 (batch dim included)
+ AddTestWeights<int32>("weights", {1}, {-5});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "Modifying batch dimension is not supported for ExpandDims, at "
+ "my_expanddims");
+ }
+ {
+ // Axis > rank(input), should fail.
+ Reset();
+ AddTestTensor("input", {1, 2, 3});
+ // Input is rank 4 (batch dim included)
+ AddTestWeights<int32>("weights", {1}, {5});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "Axis for ExpandDims is invalid, must be in the range "
+ "[-rank(input) - 1, rank(input)], at my_expanddims");
+ }
+ {
+ // Axis < -rank(input)-1, should fail.
+ Reset();
+ AddTestTensor("input", {1, 2, 3});
+ // Input is rank 4 (batch dim included)
+ AddTestWeights<int32>("weights", {1}, {-6});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "Axis for ExpandDims is invalid, must be in the range "
+ "[-rank(input) - 1, rank(input)], at my_expanddims");
+ }
+
+ struct TestParams {
+ TestParams(const std::vector<int>& input_dims, int axis,
+ const std::vector<int>& expected_output_dims)
+ : input_dims(input_dims),
+ axis(axis),
+ expected_output_dims(expected_output_dims) {}
+ std::vector<int> input_dims;
+ int axis;
+ std::vector<int> expected_output_dims;
+ };
+
+ // Ok.
+ const int kExpandDimsOKCases = 8;
+ TestParams ok_params[kExpandDimsOKCases] = {
+ TestParams{{2, 3}, 1, {1, 2, 3}}, TestParams{{2, 3}, -3, {1, 2, 3}},
+ TestParams{{2, 3}, 3, {2, 3, 1}}, TestParams{{2, 3}, -1, {2, 3, 1}},
+ TestParams{{2, 3}, 2, {2, 1, 3}}, TestParams{{2, 3}, -2, {2, 1, 3}},
+ TestParams{{6}, 1, {1, 6}}, TestParams{{6}, -1, {6, 1}},
+ };
+ for (int i = 0; i < kExpandDimsOKCases; ++i) {
+ Reset();
+ AddTestTensor("input", ok_params[i].input_dims);
+ AddTestWeights<int32>("weights", {1}, {ok_params[i].axis});
+ RunValidationAndConversion(node_def);
+ TRT_TensorOrWeights output;
+ TF_EXPECT_OK(GetTensorOrWeights("my_expanddims", &output));
+ EXPECT_TRUE(output.is_tensor());
+ ExpectTrtDimsEqualsArray(ok_params[i].expected_output_dims,
+ output.tensor()->getDimensions());
+
+ std::vector<float> output_data(6);
+ BuildAndRun<float>({{"input", {1, 2, 3, 4, 5, 6}}}, "my_expanddims",
+ &output_data);
+ EXPECT_THAT(output_data, ElementsAre(1, 2, 3, 4, 5, 6));
+ }
+}
+
+TEST_F(OpConverterTest, ConvertSqueeze) {
+ {
+ // Input list is empty, should fail.
+ NodeDef node_def = MakeNodeDef("my_squeeze", "Squeeze", {});
+ RunValidationAndConversion(node_def, error::INVALID_ARGUMENT,
+ "One input expected for Squeeze, at my_squeeze");
+ }
+ {
+ // No attrs, should fail.
+ Reset();
+ Scope s = Scope::NewRootScope();
+ auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT);
+ auto squeeze = ops::Squeeze(s.WithOpName("my_squeeze"), input);
+ const NodeDef& node_def = squeeze.operation.node()->def();
+ AddTestTensor("input", {1, 2, 3});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "Squeeze is only implemented for explicit dims, at my_squeeze");
+ }
+
+ // Get the NodeDef for Squeeze.
+ auto get_squeeze_nodedef = [](std::vector<int> axis) -> NodeDef {
+ Scope s = Scope::NewRootScope();
+ auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT);
+ ops::Squeeze::Attrs squeeze_attrs;
+ squeeze_attrs.axis_ = gtl::ArraySlice<int>(axis);
+ auto squeeze =
+ ops::Squeeze(s.WithOpName("my_squeeze"), input, squeeze_attrs);
+ return squeeze.operation.node()->def();
+ };
+
+ {
+ // Input is weights, should fail.
+ Reset();
+ NodeDef node_def = get_squeeze_nodedef({0});
+ AddTestWeights<float>("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "Squeeze expects tensor for input, at my_squeeze");
+ }
+ {
+ // Squeeze batch dim, should fail.
+ Reset();
+ NodeDef node_def = get_squeeze_nodedef({0});
+ AddTestTensor("input", {1, 2, 3});
+ RunValidationAndConversion(node_def, error::UNIMPLEMENTED,
+ "Cannot squeeze batch dimension, at my_squeeze");
+ }
+ {
+ // Squeeze batch dim via negative axis, should fail.
+ Reset();
+ NodeDef node_def = get_squeeze_nodedef({-4});
+ AddTestTensor("input", {1, 2, 3});
+ RunValidationAndConversion(node_def, error::UNIMPLEMENTED,
+ "Cannot squeeze batch dimension, at my_squeeze");
+ }
+ {
+ // Squeeze >= rank(input), should fail.
+ Reset();
+ NodeDef node_def = get_squeeze_nodedef({4});
+ AddTestTensor("input", {1, 2, 3});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "Axis for Squeeze is invalid, must be in the range "
+ "[-rank(input), rank(input)), at my_squeeze");
+ }
+ {
+ // Squeeze < -rank(input), should fail.
+ Reset();
+ NodeDef node_def = get_squeeze_nodedef({-5});
+ AddTestTensor("input", {1, 2, 3});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "Axis for Squeeze is invalid, must be in the range "
+ "[-rank(input), rank(input)), at my_squeeze");
+ }
+
+ struct TestParams {
+ TestParams(const std::vector<int>& input_dims, const std::vector<int>& axis,
+ const std::vector<int>& expected_output_dims)
+ : input_dims(input_dims),
+ axis(axis),
+ expected_output_dims(expected_output_dims) {}
+ std::vector<int> input_dims;
+ std::vector<int> axis;
+ std::vector<int> expected_output_dims;
+ };
+
+ // Ok.
+ const int kSqueezeOKCases = 10;
+ TestParams ok_params[kSqueezeOKCases] = {
+ TestParams{{1, 2, 3}, {1}, {2, 3}},
+ TestParams{{1, 2, 3}, {-3}, {2, 3}},
+ TestParams{{2, 3, 1}, {3}, {2, 3}},
+ TestParams{{2, 3, 1}, {-1}, {2, 3}},
+ TestParams{{1, 2, 1, 3, 1}, {1, 3, 5}, {2, 3}},
+ TestParams{{1, 2, 1, 3, 1}, {3, 1, 5}, {2, 3}},
+ TestParams{{1, 2, 1, 3, 1}, {-1, -3, -5}, {2, 3}},
+ TestParams{{1, 2, 1, 3, 1}, {1, -3, 5}, {2, 3}},
+ TestParams{{1, 6}, {1}, {6}},
+ TestParams{{6, 1}, {2}, {6}},
+ };
+ for (int i = 0; i < kSqueezeOKCases; ++i) {
+ Reset();
+ NodeDef node_def = get_squeeze_nodedef(ok_params[i].axis);
+ AddTestTensor("input", ok_params[i].input_dims);
+ RunValidationAndConversion(node_def);
+ TRT_TensorOrWeights output;
+ TF_EXPECT_OK(GetTensorOrWeights("my_squeeze", &output));
+ EXPECT_TRUE(output.is_tensor());
+ ExpectTrtDimsEqualsArray(ok_params[i].expected_output_dims,
+ output.tensor()->getDimensions());
+
+ std::vector<float> output_data(6);
+ BuildAndRun<float>({{"input", {1, 2, 3, 4, 5, 6}}}, "my_squeeze",
+ &output_data);
+ EXPECT_THAT(output_data, ElementsAre(1, 2, 3, 4, 5, 6));
+ }
+}
+
+TEST_F(OpConverterTest, ConvertStridedSlice) {
+ {
+ // Input list is empty, should fail.
+ NodeDef node_def = MakeNodeDef("my_strided_slice", "StridedSlice", {});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "StridedSlice expects 4 inputs, at my_strided_slice");
+ }
+
+ // Get nodedef for StridedSlice layer.
+ auto get_strided_slice_nodedef =
+ [](int begin_mask = 0, int end_mask = 0, int ellipsis_mask = 0,
+ int new_axis_mask = 0, int shrink_axis_mask = 0) -> NodeDef {
+ Scope s = Scope::NewRootScope();
+ auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT);
+ auto begin = ops::Placeholder(s.WithOpName("begin"), DT_INT32);
+ auto end = ops::Placeholder(s.WithOpName("end"), DT_INT32);
+ auto strides = ops::Placeholder(s.WithOpName("strides"), DT_INT32);
+ ops::StridedSlice::Attrs attrs = ops::StridedSlice::Attrs()
+ .BeginMask(begin_mask)
+ .EndMask(end_mask)
+ .EllipsisMask(ellipsis_mask)
+ .NewAxisMask(new_axis_mask)
+ .ShrinkAxisMask(shrink_axis_mask);
+ auto strided_slice = ops::StridedSlice(s.WithOpName("my_strided_slice"),
+ input, begin, end, strides, attrs);
+ return strided_slice.operation.node()->def();
+ };
+
+ {
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestWeights<int32>("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6});
+ AddTestWeights<int32>("begin", {4}, {0, 0, 0, 0});
+ AddTestWeights<int32>("end", {4}, {1, 1, 2, 3});
+ AddTestWeights<int32>("strides", {4}, {1, 1, 1, 1});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "StridedSlice is only implemented for tensors, at my_strided_slice");
+ }
+ {
+ // Begin, end, strides are tensors, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestTensor("begin", {4});
+ AddTestTensor("end", {4});
+ AddTestTensor("strides", {4});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "StridedSlice expects weights for begin, end, and strides, at "
+ "my_strided_slice");
+ }
+ {
+ // Non-zero ellipsis_mask, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef(
+ /*begin_mask=*/0, /*end_mask=*/0, /*ellipsis_mask=*/2,
+ /*new_axis_mask=*/0, /*shrink_axis_mask=*/0);
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("begin", {4}, {0, 0, 0, 0});
+ AddTestWeights<int32>("end", {4}, {1, 1, 2, 3});
+ AddTestWeights<int32>("strides", {4}, {1, 1, 1, 1});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "ellipsis_mask is not supported for StridedSlice, at "
+ "my_strided_slice");
+ }
+ {
+ // Modify batch dim, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("begin", {4}, {0, 0, 0, 0});
+ AddTestWeights<int32>("end", {4}, {0, 1, 2, 3});
+ AddTestWeights<int32>("strides", {4}, {1, 1, 1, 1});
+ RunValidationAndConversion(
+ node_def, error::UNIMPLEMENTED,
+ "StridedSlice can't modify batch dim, at my_strided_slice");
+ }
+ {
+ // Stride is not 1, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("begin", {4}, {0, 0, 0, 0});
+ AddTestWeights<int32>("end", {4}, {1, 1, 2, 3});
+ AddTestWeights<int32>("strides", {4}, {1, 2, -1, 3});
+ RunValidationAndConversion(node_def, error::UNIMPLEMENTED,
+ "StridedSlice is only implemented for stride of "
+ "1, at my_strided_slice");
+ }
+ {
+ // Begin out of bounds, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("begin", {4}, {1, 2, 3, 4});
+ AddTestWeights<int32>("end", {4}, {0, 1, 2, 3});
+ AddTestWeights<int32>("strides", {4}, {1, 1, 1, 1});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "begin value of 2 for StridedSlice is invalid, must be in the range "
+ "[-dim_size(i), dim_size(i)], at my_strided_slice");
+ }
+ {
+ // End out of bounds, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("begin", {4}, {0, 0, 0, 0});
+ AddTestWeights<int32>("end", {4}, {1, 2, 3, 4});
+ AddTestWeights<int32>("strides", {4}, {1, 1, 1, 1});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "end value of 2 for StridedSlice is invalid, must be in the range "
+ "[-dim_size(i), dim_size(i)], at my_strided_slice");
+ }
+ {
+ // Size of sliced dim is negative, should fail.
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef();
+ AddTestTensor("input", {1, 2, 3});
+ AddTestWeights<int32>("begin", {4}, {0, 0, 2, 0});
+ AddTestWeights<int32>("end", {4}, {1, 1, 0, 3});
+ AddTestWeights<int32>("strides", {4}, {1, 1, 1, 1});
+ RunValidationAndConversion(
+ node_def, error::INVALID_ARGUMENT,
+ "New size of sliced dimension is negative, at my_strided_slice");
+ }
+
+ struct TestParams {
+ TestParams(const std::vector<int>& input_dims,
+ const std::vector<int>& expected_output_dims,
+ const std::vector<int>& begin, const std::vector<int>& end,
+ const std::vector<int>& begin_mask,
+ const std::vector<int>& end_mask,
+ const std::vector<int>& expected_output)
+ : input_dims(input_dims),
+ expected_output_dims(expected_output_dims),
+ begin(begin),
+ end(end),
+ expected_output(expected_output) {
+ // Masks are provided in terms of vectors for readability. Convert them to
+ // binary here.
+ this->begin_mask = 0;
+ for (int i = 0; i < begin_mask.size(); i++) {
+ if (begin_mask[i]) this->begin_mask |= (1 << i);
+ }
+ this->end_mask = 0;
+ for (int i = 0; i < end_mask.size(); i++) {
+ if (end_mask[i]) this->end_mask |= (1 << i);
+ }
+ }
+
+ std::vector<int> input_dims;
+ std::vector<int> expected_output_dims;
+ std::vector<int> begin;
+ std::vector<int> end;
+ int begin_mask;
+ int end_mask;
+ std::vector<int> expected_output;
+ };
+
+ // Ok.
+ const int kStridedSliceOKCases = 18;
+ TestParams ok_params[kStridedSliceOKCases] = {
+ // 2D Crop.
+ TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 2},
+ /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 0, 1, 2},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 0, 0},
+ /*expected_output=*/{1, 2}},
+ TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 2},
+ /*begin=*/{0, 0, 1, 1}, /*end=*/{0, 0, 0, 0},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 1},
+ /*expected_output=*/{5, 6}},
+ TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 2},
+ /*begin=*/{0, 0, 1, 1}, /*end=*/{0, 1, 2, 3},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 0, 0},
+ /*expected_output=*/{5, 6}},
+ // 2D Crop, with transpose.
+ TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 2, 1},
+ /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 1, 2, 1},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0},
+ /*expected_output=*/{1, 2}},
+ TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 2, 1},
+ /*begin=*/{0, 1, 1, 0}, /*end=*/{0, 2, 3, 1},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0},
+ /*expected_output=*/{5, 6}},
+ TestParams{/*input_dims=*/{2, 1, 3}, /*expected_output_dims=*/{1, 1, 2},
+ /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 1, 1, 2},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0},
+ /*expected_output=*/{1, 2}},
+ TestParams{/*input_dims=*/{2, 1, 3}, /*expected_output_dims=*/{1, 1, 2},
+ /*begin=*/{0, 1, 0, 1}, /*end=*/{0, 2, 1, 3},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0},
+ /*expected_output=*/{5, 6}},
+ // 2D Crop, with reshape.
+ TestParams{/*input_dims=*/{2, 3}, /*expected_output_dims=*/{1, 2},
+ /*begin=*/{0, 0, 0}, /*end=*/{0, 1, 2},
+ /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 0},
+ /*expected_output=*/{1, 2}},
+ TestParams{/*input_dims=*/{2, 3}, /*expected_output_dims=*/{1, 2},
+ /*begin=*/{0, 1, 1}, /*end=*/{0, 0, 0},
+ /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 1, 1},
+ /*expected_output=*/{5, 6}},
+ // 1D Crop.
+ TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 2, 2},
+ /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 0, 0, 2},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 0},
+ /*expected_output=*/{1, 2, 4, 5}},
+ TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 3},
+ /*begin=*/{0, 0, 1, 0}, /*end=*/{0, 0, 0, 0},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 1},
+ /*expected_output=*/{4, 5, 6}},
+ // 1D Crop, with transpose.
+ TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 3, 1},
+ /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 1, 0, 0},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 1, 1},
+ /*expected_output=*/{1, 2, 3}},
+ TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 3, 1},
+ /*begin=*/{0, 1, 0, 0}, /*end=*/{0, 0, 0, 0},
+ /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 1},
+ /*expected_output=*/{4, 5, 6}},
+ // 1D Crop, with reshape.
+ TestParams{/*input_dims=*/{6}, /*expected_output_dims=*/{3},
+ /*begin=*/{0, 0}, /*end=*/{0, 3},
+ /*begin_mask=*/{0, 0}, /*end_mask=*/{1, 0},
+ /*expected_output=*/{1, 2, 3}},
+ TestParams{/*input_dims=*/{1, 6}, /*expected_output_dims=*/{1, 3},
+ /*begin=*/{0, 0, 2}, /*end=*/{0, 0, 5},
+ /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 1, 0},
+ /*expected_output=*/{3, 4, 5}},
+ TestParams{/*input_dims=*/{6, 1}, /*expected_output_dims=*/{3, 1},
+ /*begin=*/{0, 2, 0}, /*end=*/{0, 5, 0},
+ /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 1},
+ /*expected_output=*/{3, 4, 5}},
+ // Negative axis.
+ TestParams{/*input_dims=*/{6, 1}, /*expected_output_dims=*/{3, 1},
+ /*begin=*/{0, -6, 0}, /*end=*/{0, -3, 0},
+ /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 1},
+ /*expected_output=*/{1, 2, 3}},
+ TestParams{/*input_dims=*/{6, 1}, /*expected_output_dims=*/{5, 1},
+ /*begin=*/{0, 0, 0}, /*end=*/{0, -1, 0},
+ /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 1},
+ /*expected_output=*/{1, 2, 3, 4, 5}},
+ };
+
+ for (int i = 0; i < kStridedSliceOKCases; i++) {
+ Reset();
+ NodeDef node_def = get_strided_slice_nodedef(ok_params[i].begin_mask,
+ ok_params[i].end_mask);
+ AddTestTensor("input", ok_params[i].input_dims);
+ AddTestWeights<int32>("begin",
+ {static_cast<int>(ok_params[i].begin.size())},
+ ok_params[i].begin);
+ AddTestWeights<int32>("end", {static_cast<int>(ok_params[i].end.size())},
+ ok_params[i].end);
+ std::vector<int> strides(ok_params[i].input_dims.size(), 1);
+ AddTestWeights<int32>("strides", {static_cast<int>(strides.size())},
+ strides);
+ RunValidationAndConversion(node_def);
+
+ TRT_TensorOrWeights output;
+ TF_EXPECT_OK(GetTensorOrWeights("my_strided_slice", &output));
+ std::vector<float> output_data(ok_params[i].expected_output.size());
+ BuildAndRun<float>({{"input", {1, 2, 3, 4, 5, 6}}}, "my_strided_slice",
+ &output_data);
+ EXPECT_THAT(output_data, ElementsAreArray(ok_params[i].expected_output));
+ }
+}
+
} // namespace convert
} // namespace tensorrt
} // namespace tensorflow
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
index c1688d4..d57f230 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -226,8 +226,9 @@
tensorflow::tensorrt::convert::ConversionParams cp;
if (use_calibration_ && precision_mode_ != INT8MODE) {
- LOG(ERROR) << "Calibration with FP32 or FP16 is not implemented. "
- << "Falling back to use_calibration = False.";
+ VLOG(1) << "Calibration with FP32 or FP16 is not implemented. "
+ << "Falling back to use_calibration = False."
+ << "Note that the default value of use_calibration is True.";
use_calibration_ = false;
}
diff --git a/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py b/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py
index 31cbef8..e7d6ec4 100644
--- a/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py
+++ b/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py
@@ -191,7 +191,7 @@
batch_size=batch_size,
num_parallel_calls=8))
dataset = dataset.repeat(count=1)
- iterator = data.make_one_shot_iterator(dataset)
+ iterator = dataset.make_one_shot_iterator()
features, labels = iterator.get_next()
return features, labels
@@ -205,7 +205,7 @@
batch_size=batch_size,
num_parallel_calls=8))
dataset = dataset.repeat(count=num_epochs)
- iterator = data.make_one_shot_iterator(dataset)
+ iterator = dataset.make_one_shot_iterator()
features, labels = iterator.get_next()
return features, labels
diff --git a/tensorflow/contrib/tensorrt/test/rank_two_test.py b/tensorflow/contrib/tensorrt/test/rank_two_test.py
index 0cd733d..563232f 100644
--- a/tensorflow/contrib/tensorrt/test/rank_two_test.py
+++ b/tensorflow/contrib/tensorrt/test/rank_two_test.py
@@ -51,8 +51,10 @@
c = constant_op.constant(3.0, name="c%d_3" % i)
q = math_ops.add(q, c, name="add%d_3" % i)
if i == 0:
+ axis = constant_op.constant(-1, dtype=dtypes.int32, name="axis")
for j in range(2):
- q = array_ops.expand_dims(q, -1, name="expand%d_%d" % (i, j))
+ q = array_ops.expand_dims(q, axis, name="expand%d_%d" % (i, j))
+ q = self.trt_incompatible_op(q)
q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i)
outputs.append(q)
# Combine both paths
@@ -70,7 +72,7 @@
return {
"TRTEngineOp_0": [
"add0_1", "add0_2", "add0_3", "c0_1", "c0_2", "c0_3", "abs0_1",
- "abs0_2"
+ "abs0_2", "expand0_0", "expand0_1", "axis"
],
"TRTEngineOp_1": [
"add", "add1_1", "add1_2", "add1_3", "c1_1", "c1_2", "c1_3",
diff --git a/tensorflow/contrib/tensorrt/test/unary_test.py b/tensorflow/contrib/tensorrt/test/unary_test.py
index 9fc50e0..b6e5e32 100644
--- a/tensorflow/contrib/tensorrt/test/unary_test.py
+++ b/tensorflow/contrib/tensorrt/test/unary_test.py
@@ -106,10 +106,7 @@
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
- return [
- "TRTEngineOp_0", "TRTEngineOp_1", "TRTEngineOp_2", "TRTEngineOp_3",
- "TRTEngineOp_4"
- ]
+ return ["TRTEngineOp_0"]
if __name__ == "__main__":
diff --git a/tensorflow/contrib/tfprof/README.md b/tensorflow/contrib/tfprof/README.md
index b29d1ac..f40e76f 100644
--- a/tensorflow/contrib/tfprof/README.md
+++ b/tensorflow/contrib/tfprof/README.md
@@ -1,7 +1,5 @@
# tfprof: TensorFlow Profiler and Beyond
-<h1>Please use `tf.profiler.xxx` instead of `tf.contrib.tfprof.xxx`</h1>
-
<h1>Full Document in
<a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/README.md">tensorflow/core/profiler/README.md</a><h1>
diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD
index 05d2ebd..007aeae 100644
--- a/tensorflow/contrib/tpu/BUILD
+++ b/tensorflow/contrib/tpu/BUILD
@@ -79,6 +79,7 @@
"//tensorflow/python:init_ops",
"//tensorflow/python:math_ops",
"//tensorflow/python:platform",
+ "//tensorflow/python:session",
"//tensorflow/python:state_ops",
"//tensorflow/python:summary",
"//tensorflow/python:summary_ops_v2",
@@ -101,6 +102,7 @@
"replication_ops",
"tpu_configuration_ops",
"tpu_embedding_ops",
+ "tpu_ordinal_selector_op",
],
deps = [
"//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc",
@@ -152,6 +154,13 @@
],
)
+tf_gen_op_wrapper_py(
+ name = "tpu_ordinal_selector_op",
+ deps = [
+ ":tpu_ordinal_selector_op_op_lib",
+ ],
+)
+
py_library(
name = "profiler",
srcs = ["python/profiler/__init__.py"],
diff --git a/tensorflow/contrib/tpu/ops/tpu_ordinal_selector_op.cc b/tensorflow/contrib/tpu/ops/tpu_ordinal_selector_op.cc
new file mode 100644
index 0000000..54e6b20
--- /dev/null
+++ b/tensorflow/contrib/tpu/ops/tpu_ordinal_selector_op.cc
@@ -0,0 +1,39 @@
+/* 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/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+namespace tensorflow {
+
+REGISTER_OP("TPUOrdinalSelector")
+ .Output("device_ordinals: int32")
+ .SetIsStateful()
+ .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+ c->set_output(0,
+ c->Vector(shape_inference::InferenceContext::kUnknownDim));
+ return Status::OK();
+ })
+ .Doc(R"doc(
+A TPU core selector Op.
+
+This Op produces a set of TPU cores (for warm-up) or a single TPU core
+(for regular inference) to execute the TPU program on. The output is
+consumed by TPUPartitionedCall.
+
+device_ordinals: A vector 1 or more TPU cores.
+)doc");
+
+} // namespace tensorflow
diff --git a/tensorflow/contrib/tpu/python/tpu/keras_support.py b/tensorflow/contrib/tpu/python/tpu/keras_support.py
index 4ce1945..cf9672f 100644
--- a/tensorflow/contrib/tpu/python/tpu/keras_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/keras_support.py
@@ -2069,6 +2069,8 @@
# tpu_model may not be compiled, e.g., loading weights and then predict.
return
for k, v in six.iteritems(cpu_optimizer_config):
+ if k == 'name':
+ continue
opt_var = getattr(self._tpu_model.optimizer, k)
if isinstance(opt_var, variables.Variable):
logging.info('CPU -> TPU %s: %s {%s}', k, v, K.get_value(opt_var))
@@ -2097,6 +2099,8 @@
self._cpu_model.set_weights(tpu_weights)
for k, v in six.iteritems(tpu_optimizer_config):
logging.info('TPU -> CPU %s: %s', k, v)
+ if k == 'name':
+ continue
opt_var = getattr(self.cpu_optimizer, k)
if isinstance(opt_var, variables.Variable):
K.get_session().run(opt_var.assign(v))
diff --git a/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py b/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py
index 8b0b240..de42562 100644
--- a/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py
+++ b/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py
@@ -69,6 +69,7 @@
def __init__(self, name, variables):
self._name = name
self._primary_var = variables[0]
+ self._common_name = self._primary_var.name.split(":")[0]
self._vars = variables
self._cached_value = None
self._dtype = variables[0].dtype
diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py
index a952754..3e46382 100644
--- a/tensorflow/contrib/tpu/python/tpu/session_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/session_support.py
@@ -43,12 +43,19 @@
pass
+def _clone_session(session, graph=None):
+ return session_lib.Session(
+ target=session.sess_str,
+ config=session._config, # pylint: disable=protected-access
+ graph=graph if graph else session.graph)
+
+
def _make_heartbeat_op(session, device, request_ph):
"""Return a heartbeat op or None if heartbeats are not supported by device."""
try:
# Test if we can connect in a isolated graph + session
with ops.Graph().as_default():
- with session_lib.Session(target=session.sess_str) as temp_session:
+ with _clone_session(session) as temp_session:
with ops.device(device):
heartbeat_op = tpu_ops.worker_heartbeat('')
options = config_pb2.RunOptions(timeout_in_ms=5000)
@@ -220,6 +227,7 @@
self.ping_interval = ping_interval
self.shutdown_timeout = shutdown_timeout
self.daemon = True
+ self._config = session._config # pylint: disable=protected-access
self._target = session.sess_str
self._running = False
self._devices = devices
@@ -234,6 +242,7 @@
self._session = session_lib.Session(
target=self._target,
graph=self._graph,
+ config=self._config,
)
if self._devices is None:
@@ -334,8 +343,7 @@
with self._graph.as_default():
logging.info('Installing graceful shutdown hook.')
- self._session = session_lib.Session(
- target=training_session.sess_str, graph=self._graph)
+ self._session = _clone_session(training_session, self._graph)
self._workers = WorkerHeartbeatManager.from_devices(
self._session, all_worker_devices(self._session))
self._heartbeat_supported = self._workers.num_workers() > 0
diff --git a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py
index 70baea2..a1494e3 100644
--- a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py
+++ b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py
@@ -21,44 +21,56 @@
import os
import os.path
import re
+import sys
from tensorflow.contrib.tpu.python.ops import tpu_ops
from tensorflow.contrib.tpu.python.tpu import tpu
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_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import logging_ops
from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
_TRACER_LOG_PREFIX = ' [>>>TT>>>]'
_DEVICE_TYPE_TPU = 'tpu'
_DEVICE_TYPE_CPU = 'cpu'
-_GLOBAL_STEP_OP_NAME = 'GLOBAL-STEP'
_TRACE_MODE_NAN_INF = 'nan-inf'
_TRACE_MODE_PART_TENSOR = 'part-tensor'
_TRACE_MODE_PART_TENSOR_SIZE = 3
_TRACE_MODE_FULL_TENSOR = 'full-tensor'
-_RECORD_OUTSIDE_OP_RANGE = 'not-traced-outside-op-range'
-_RECORD_SHOULD_NOT_TRACE = 'not-traced-should-not-trace'
-_RECORD_FILTERED_OUT = 'not-traced-filtered-out'
-_RECORD_SCALAR = 'not-traced-scalar'
-_RECORD_DYNAMIC_SHAPE = 'not-traced-dynamic-shape'
-_RECORD_GET_TRACED = 'get-traced'
+_TRACE_MODE_NORM = 'norm'
+_TRACE_MODE_MAX_ABS = 'max-abs'
+_REASON_OUTSIDE_OP_RANGE = 'not-traced-outside-op-range'
+_REASON_UNSAFE_OP = 'not-traced-unsafe-op'
+_REASON_UNSAFE_SCALAR = 'not-traced-unsafe-scalar'
+_REASON_LESS_INTERESTING_OP = 'not-traced-less-interesting-op'
+_REASON_DEVICE_MISMATCH = 'not-traced-device-mismatch'
+_REASON_DYNAMIC_SHAPE = 'not-traced-dynamic-shape'
+_REASON_SCALAR_GET_TRACED = 'traced-scalar'
+_REASON_TENSOR_GET_TRACED = 'traced-tensor'
+_REASON_USER_INCLUDED = 'traced-user-included'
+_REASON_USER_EXCLUDED = 'not-traced-user-excluded'
+_REASON_NON_NUMERIC_TENSOR = 'not-traced-non-numeric-tensor'
_MARKER_SECTION_BEGIN = '!!!!!!! section-begin:'
_MARKER_SECTION_END = '!!!!!!! section-end:'
_SECTION_NAME_CONFIG = 'configuration'
_SECTION_NAME_REASON = 'reason'
_SECTION_NAME_OP_LIST = 'op-list'
+_SECTION_NAME_TENSOR_LIST = 'tensor-list'
_SECTION_NAME_GRAPH = 'graph'
_FIELD_NAME_VERSION = 'version:'
_FIELD_NAME_DEVICE = 'device:'
_FIELD_NAME_TRACE_MODE = 'trace-mode:'
_FIELD_NAME_NUM_REPLICAS = 'num-replicas:'
_FIELD_NAME_NUM_OPS = 'number-of-ops:'
+_FIELD_NAME_NUM_TENSORS = 'number-of-tensors:'
_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED = 'topological-sort-succeed:'
_FLAGS_ENV_VAR = 'TENSOR_TRACER_FLAGS'
_FLAG_SINGLE_QUOTE_PAT = re.compile(r"\s*--([^=]+)='([^']*)'")
@@ -66,13 +78,72 @@
_FLAG_NO_QUOTE_PAT = re.compile(r'\s*--([^=]+)=(\S*)')
_FLAG_NAME_ENABLE = 'enable'
_FLAG_NAME_TRACE_MODE = 'trace_mode'
-_FLAG_NAME_INTERESTING_OPS = 'interesting_ops'
+_FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS = 'include_less_interesting_ops'
+_FLAG_NAME_EXCLUDED_OPNAMES = 'excluded_opnames'
+_FLAG_NAME_EXCLUDED_OPTYPES = 'excluded_optypes'
+_FLAG_NAME_INCLUDED_OPNAMES = 'included_opnames'
+_FLAG_NAME_INCLUDED_OPTYPES = 'included_optypes'
_FLAG_NAME_TRACE_FILE = 'trace_file_path'
+_FLAG_NAME_REPORT_FILE = 'report_file_path'
_FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR = 'use_test_undeclared_outputs_dir'
_FLAG_NAME_OP_RANGE = 'op_range'
_OP_RANGE_PAT = re.compile(r'(\d+):(\d+)')
_OUTPUT_STREAM_ESCAPE = 'file://'
_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR = 'TEST_UNDECLARED_OUTPUTS_DIR'
+_TENSOR_TRACER_COLLECTION = 'tensor_tracer_variables'
+_TENSOR_TRACER_CHECKPOINT = 'tensor_tracer_checkpoint'
+
+
+def tensor_checkpoint(tensor, checkpoint_name):
+ """Adds a checkpoint with the given checkpoint name for the given tensor.
+
+ The tensor will be added to the list of tensors that will be traced by the
+ tensor tracer.
+
+ Args:
+ tensor: the tensor object for which the tracing is requested.
+ checkpoint_name: a string name for the checkpoint. This name has to be a
+ unique name if used within model comparison. The tensors that have the same
+ checkpoint identifier is compared in model comparison.
+ Returns:
+ The provided tensor.
+ """
+
+ tensor.graph.get_collection(_TENSOR_TRACER_COLLECTION)
+ tensor.graph.add_to_collection(_TENSOR_TRACER_COLLECTION,
+ (tensor, checkpoint_name))
+ return tensor
+
+
+def keras_layer_checkpoint(layer, checkpoint_name):
+ """An interface for adding the tensor outputs of a keras layer.
+
+ Encapsulates tensor_checkpoint.
+
+ Args:
+ layer: A keras layer.
+ checkpoint_name: a string name for the checkpoint. This name has to be a
+ unique name if used within model comparison. The tensors that have the same
+ checkpoint identifier is compared in model comparison.
+
+ Returns:
+ The provided layer.
+ """
+ try:
+ outputs = layer.output
+ if tensor_util.is_tensor(outputs):
+ tensor_checkpoint(outputs, '%s' % (checkpoint_name))
+ else:
+ idx = 0
+ for output_tensor in outputs:
+ if tensor_util.is_tensor(outputs):
+ tensor_checkpoint(output_tensor, '%s_%d' % (checkpoint_name, idx))
+ idx += 1
+ except AttributeError:
+ pass
+ except RuntimeError:
+ pass
+ return layer
class TensorTracer(object):
@@ -106,6 +177,34 @@
return match
@staticmethod
+ def validate_flag_names():
+ """Validates if the TensorTrace flags passed are valid."""
+ valid_flag_names = [_FLAG_NAME_ENABLE, _FLAG_NAME_TRACE_MODE,
+ _FLAG_NAME_EXCLUDED_OPNAMES,
+ _FLAG_NAME_EXCLUDED_OPTYPES,
+ _FLAG_NAME_INCLUDED_OPNAMES,
+ _FLAG_NAME_INCLUDED_OPTYPES,
+ _FLAG_NAME_TRACE_FILE, _FLAG_NAME_REPORT_FILE,
+ _FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR,
+ _FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS,
+ _FLAG_NAME_OP_RANGE]
+ tensor_tracer_flags = os.environ.get(_FLAGS_ENV_VAR)
+ if not tensor_tracer_flags:
+ return
+ pos = 0
+ while True:
+ match = TensorTracer._match_next_flag(tensor_tracer_flags, pos)
+ if not match:
+ break
+ flag_name = match.group(1)
+ if flag_name not in valid_flag_names:
+ raise ValueError(
+ 'The flag name "%s" passed via the environment variable "%s" '
+ 'is invalid. Valid flag names are:'
+ '\n%s'%(flag_name, _FLAGS_ENV_VAR, valid_flag_names))
+ pos = match.end()
+
+ @staticmethod
def print_flag_values():
"""Prints all TensorTracer flags passed via environment variables."""
@@ -147,6 +246,20 @@
return ''
@staticmethod
+ def flag_value_to_re_list(flag_name):
+ """Converts list of strings to compiled RE."""
+
+ re_list = []
+ flag_value = TensorTracer.get_flag_value(flag_name)
+ if not flag_value:
+ return re_list
+ list_of_values = flag_value.split()
+ for v in list_of_values:
+ r = re.compile(v)
+ re_list.append(r)
+ return re_list
+
+ @staticmethod
def is_enabled():
"""Returns True if TensorTracer is enabled."""
@@ -186,29 +299,67 @@
"""Checks if the given trace mode is valid."""
valid_trace_modes = [_TRACE_MODE_NAN_INF, _TRACE_MODE_PART_TENSOR,
- _TRACE_MODE_FULL_TENSOR]
+ _TRACE_MODE_FULL_TENSOR, _TRACE_MODE_NORM,
+ _TRACE_MODE_MAX_ABS]
if trace_mode not in valid_trace_modes:
raise ValueError('Invalid trace mode "%s" given to the Tensor_Tracer.'
'Valid trace modes are: %s'%(trace_mode,
valid_trace_modes))
@staticmethod
- def should_trace(device_type, op):
- """Returns True if the given Op should be traced."""
+ def unsafe_op(op):
+ """Returns True if this op is not safe to be traced."""
- if device_type != _DEVICE_TYPE_TPU:
- raise ValueError('Non TPU device type is not supported')
if control_flow_util.IsInCond(op):
+ return True
+ # Reasons for not including following op types:
+ # Assign: cause incorrect result with CPU tracing.
+ # others: compilation problems.
+ if op.type in ['Assign', 'Pack', 'Shape', 'Reshape', 'ArgMin', 'ArgMax']:
+ return True
+ return False
+
+ @staticmethod
+ def device_mismatch(device_type, op):
+ if device_type == _DEVICE_TYPE_TPU:
+ # pylint: disable=protected-access
+ return tpu._TPU_REPLICATE_ATTR not in op.node_def.attr
+ # pylint: enable=protected-access
+ return False
+
+ @staticmethod
+ def unsafe_scalar_trace(op):
+ """Return true if scalar output tensor from Op is not safe to be traced."""
+
+ # Tracing the following causes cycle in the graph on TPU.
+ if op.type in ['LoopCond', 'Enter', 'Merge', 'Const',
+ 'Switch', 'Less', 'ReadVariableOp']:
+ return True
+ # Tracing the following will cause casting-issue
+ # with the norm tracing mode or other compilation issues on CPU.
+ if op.type in ['VarHandleOp', 'IteratorToStringHandle',
+ 'IteratorGetNext', 'OneShotIterator',
+ 'IteratorV2', 'MakeIterator',
+ 'BatchDatasetV2', 'MapDataset',
+ 'FixedLengthRecordDataset', 'TakeDataset', 'ZipDataset',
+ 'Placeholder', 'PlaceholderWithDefault', 'StridedSlice']:
+ return True
+ return False
+
+ @staticmethod
+ def less_interesting_op(op):
+ """Returns True if the given Op is not an interesting one to be traced."""
+
+ include_less_interesting = TensorTracer.get_flag_value(
+ _FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS)
+ if include_less_interesting:
return False
- if op.type in ['Reshape', 'ArgMin', 'ArgMax']:
- return False
- # pylint: disable=protected-access
- return tpu._TPU_REPLICATE_ATTR in op.node_def.attr
- # pylint: enable=protected-access
+ return op.type in ['Const', 'Identity', 'Cast', 'Shape']
@staticmethod
def reason(op_idx, details):
- """Returns why the Op at op_idx is traced or not."""
+ """Returns reason why the Op at op_idx is traced or not."""
+
return '%d %s'%(op_idx, details)
@staticmethod
@@ -274,6 +425,33 @@
assert len(unsorted_ops) == len(sorted_ops)
return (True, sorted_ops)
+ @staticmethod
+ def _make_op_and_tensor_maps(op_list):
+ """Creates various maps and lists from op_list.
+
+ Args:
+ op_list: a list of Ops
+
+ Returns:
+ opname_idx_map: a map from Op's name to its index in op_list.
+ tensor_list: a list of output tensors of the Ops in op_list.
+ tensorname_idx_map: a map from output tensor name to its index
+ in tensor_list.
+ """
+
+ opname_idx_map = {}
+ tensor_list = []
+ tensorname_idx_map = {}
+ for op_id, op in enumerate(op_list):
+ if op.name in opname_idx_map:
+ raise ValueError('Duplicated Op name: %s'%op.name)
+ opname_idx_map[op.name] = op_id
+ for output_tensor in op.outputs:
+ if output_tensor.name not in tensorname_idx_map:
+ tensor_list.append(output_tensor)
+ tensorname_idx_map[output_tensor.name] = len(tensor_list)-1
+ return (opname_idx_map, tensor_list, tensorname_idx_map)
+
def __init__(self):
"""Initializes a TensorTracer.
@@ -281,16 +459,20 @@
"""
self._version = 'use-outside-compilation'
self._device_type = None
+ TensorTracer.validate_flag_names()
self._trace_mode = TensorTracer.get_flag_value(_FLAG_NAME_TRACE_MODE)
if not self._trace_mode:
self._trace_mode = _TRACE_MODE_NAN_INF
TensorTracer.check_trace_mode(self._trace_mode)
self._part_tensor_size = _TRACE_MODE_PART_TENSOR_SIZE
self._instrument_records = {}
- interesting_ops = TensorTracer.get_flag_value(_FLAG_NAME_INTERESTING_OPS)
- self._selected_ops = interesting_ops.split()
self._set_trace_file_path()
+ self._set_report_file()
self._set_op_range()
+ self._set_excluded_opnames()
+ self._set_excluded_optypes()
+ self._set_included_opnames()
+ self._set_included_optypes()
self._num_replicas = None
self._replica_id = None
@@ -318,10 +500,7 @@
"""Sets the path of the output trace file."""
self._trace_file_path = TensorTracer.get_flag_value(_FLAG_NAME_TRACE_FILE)
- if not self._trace_file_path:
- raise ValueError('--%s is not set in the environment variable %s'
- %(_FLAG_NAME_TRACE_FILE, _FLAGS_ENV_VAR))
- elif TensorTracer.use_test_undeclared_outputs_dir():
+ if self._trace_file_path and TensorTracer.use_test_undeclared_outputs_dir():
if os.path.isabs(self._trace_file_path):
raise ValueError('If use_test_undeclared_outputs_dir is set,'
'trace_file_path cannot be an absolute path (%s)'
@@ -330,6 +509,22 @@
self._trace_file_path = os.path.join(outputs_dir,
self._trace_file_path)
+ def _set_report_file(self):
+ """Sets the path of the output report file."""
+
+ self._report_file_path = TensorTracer.get_flag_value(_FLAG_NAME_REPORT_FILE)
+ if not self._report_file_path:
+ self._report_file = None
+ return
+ try:
+ self._report_file = gfile.Open(self._report_file_path, 'w')
+ except IOError as e:
+ raise e
+
+ def _close_report_file(self):
+ if self._report_file:
+ self._report_file.close()
+
def _set_op_range(self):
"""Sets the index range of the Ops that we will consider tracing."""
@@ -350,19 +545,48 @@
return False
return self._op_range[1] < 0 or idx <= self._op_range[1]
+ def _set_excluded_opnames(self):
+ self._excluded_opname_re_list = TensorTracer.flag_value_to_re_list(
+ _FLAG_NAME_EXCLUDED_OPNAMES)
+
+ def _set_excluded_optypes(self):
+ self._excluded_optype_re_list = TensorTracer.flag_value_to_re_list(
+ _FLAG_NAME_EXCLUDED_OPTYPES)
+
+ def _set_included_opnames(self):
+ self._included_opname_re_list = TensorTracer.flag_value_to_re_list(
+ _FLAG_NAME_INCLUDED_OPNAMES)
+
+ def _set_included_optypes(self):
+ self._included_optype_re_list = TensorTracer.flag_value_to_re_list(
+ _FLAG_NAME_INCLUDED_OPTYPES)
+
+ def _is_user_included_op(self, op):
+ for opname_re in self._included_opname_re_list:
+ if opname_re.match(op.name):
+ return True
+ for optype_re in self._included_optype_re_list:
+ if optype_re.match(op.type):
+ return True
+ return False
+
+ def _is_user_excluded_op(self, op):
+ for opname_re in self._excluded_opname_re_list:
+ if opname_re.match(op.name):
+ return True
+ for optype_re in self._excluded_optype_re_list:
+ if optype_re.match(op.type):
+ return True
+ return False
+
def _write_report(self, content):
"""Writes the given content to the report."""
- logging.info('%s %s'%(_TRACER_LOG_PREFIX, content))
-
- def _is_selected_op(self, op_name):
- """Returns True if the Op with op_name is selected to be traced."""
-
- if not self._selected_ops:
- return True
- if op_name in self._selected_ops:
- return True
- return False
+ line = '%s %s'%(_TRACER_LOG_PREFIX, content)
+ if self._report_file:
+ self._report_file.write(line)
+ else:
+ logging.info(line)
def _write_config_section(self):
"""Writes the config section of the report."""
@@ -382,15 +606,42 @@
self._write_report('"%s" %s\n'%(key, self._instrument_records[key]))
self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_REASON))
- def _write_op_list_section(self, op_list):
+ def _write_op_list_section(self, op_list, tensorname_idx_map):
"""Writes the Op-list section of the report."""
self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_OP_LIST))
self._write_report('%s %d\n'%(_FIELD_NAME_NUM_OPS, len(op_list)))
for i in range(0, len(op_list)):
- self._write_report('%d "%s" %s\n'%(i, op_list[i].name, op_list[i].type))
+ op = op_list[i]
+ line = '%d "%s" %s'%(i, op.name, op.type)
+ for out_tensor in op.outputs:
+ if out_tensor.name not in tensorname_idx_map:
+ raise ValueError(
+ 'out_tensor %s is not in tensorname_idx_map'%out_tensor.name)
+ line += ' %d'%tensorname_idx_map[out_tensor.name]
+ line += '\n'
+ self._write_report(line)
self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_OP_LIST))
+ def _write_tensor_list_section(self, tensor_list, opname_idx_map):
+ """Writes the tensor-list section of the report."""
+
+ self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN,
+ _SECTION_NAME_TENSOR_LIST))
+ self._write_report('%s %d\n'%(_FIELD_NAME_NUM_TENSORS, len(tensor_list)))
+ for i in range(0, len(tensor_list)):
+ tensor = tensor_list[i]
+ line = '%d "%s"'%(i, tensor.name)
+ for consumer_op in tensor.consumers():
+ if consumer_op.name not in opname_idx_map:
+ raise ValueError(
+ 'consumer_op %s is not in opname_idx_map'%consumer_op.name)
+ line += ' %d'%opname_idx_map[consumer_op.name]
+ line += '\n'
+ self._write_report(line)
+ self._write_report('%s %s\n'%(_MARKER_SECTION_END,
+ _SECTION_NAME_TENSOR_LIST))
+
def _write_graph_section(self, succeed, sorted_or_cycle):
"""Writes the graph section of the report."""
@@ -422,7 +673,7 @@
Args:
op_name: the name of the Op that outputs the tensor to be printed.
output_idx: which output of the Op it is (0 means the first output).
- num_elements: number of elements to print.
+ num_elements: number of elements to print (-1 means print all).
tensor: the tensor needs to be returned.
output_tensor: the tensor needs to be printed.
@@ -430,10 +681,13 @@
The same tensor passed via the "tensor" argument.
"""
msg = '"%s:%d" '%(op_name, output_idx)
- output_stream = _OUTPUT_STREAM_ESCAPE + self._trace_file_path
+ if self._trace_file_path:
+ output_stream = _OUTPUT_STREAM_ESCAPE + self._trace_file_path
+ else:
+ output_stream = sys.stderr
print_op = logging_ops.print_v2(msg, array_ops.shape(output_tensor),
' @', self._replica_id,
- '\n', output_tensor,
+ '\n', output_tensor, '\n',
summarize=num_elements,
output_stream=output_stream)
with ops.control_dependencies([print_op]):
@@ -442,7 +696,8 @@
def _detect_nan_inf(tensor):
"""Trace function for detecting any NaN/Inf in the tensor."""
- if tensor.dtype.is_floating:
+ if tensor.dtype.__eq__(dtypes.bfloat16) or tensor.dtype.__eq__(
+ dtypes.float16):
# Since host can't handle bf16, always convert tensor to f32.
tensor = math_ops.cast(tensor, dtypes.float32)
output_tensor = math_ops.reduce_any(
@@ -450,12 +705,19 @@
gen_math_ops.is_inf(tensor)))
else:
output_tensor = constant_op.constant(0)
- return _print_tensor(op_name, output_idx, 1, tensor, output_tensor)
+ return _print_tensor(op_name, output_idx, -1, tensor, output_tensor)
- def _show_global_step(tensor):
- """Trace function for printing the global step count."""
+ def _show_norm(tensor):
+ tensor = math_ops.cast(tensor, dtypes.float64)
+ output_tensor = linalg_ops.norm(tensor)
+ return _print_tensor(op_name, output_idx, -1, tensor, output_tensor)
- return _print_tensor(op_name, output_idx, 1, tensor, tensor)
+ def _show_max_abs(tensor):
+ output_tensor = math_ops.cast(math_ops.reduce_max(math_ops.abs(tensor)),
+ dtypes.float64)
+ zero = constant_op.constant(0, dtypes.float64)
+ output_tensor = gen_math_ops.maximum(zero, output_tensor)
+ return _print_tensor(op_name, output_idx, -1, tensor, output_tensor)
def _show_part_tensor(tensor):
"""Trace function for printing part of the tensor."""
@@ -468,23 +730,139 @@
return _print_tensor(op_name, output_idx, -1, tensor, tensor)
- if op_name == _GLOBAL_STEP_OP_NAME:
- return _show_global_step
if self._trace_mode == _TRACE_MODE_NAN_INF:
return _detect_nan_inf
if self._trace_mode == _TRACE_MODE_PART_TENSOR:
return _show_part_tensor
if self._trace_mode == _TRACE_MODE_FULL_TENSOR:
return _show_full_tensor
+ if self._trace_mode == _TRACE_MODE_NORM:
+ return _show_norm
+ if self._trace_mode == _TRACE_MODE_MAX_ABS:
+ return _show_max_abs
raise RuntimeError('Tensor trace fun for %s is not yet implemented'
%self._trace_mode)
+ def _skip_op(self, op_id, op, user_included, user_excluded):
+ """Returns True if we should not trace Op."""
+
+ if user_included:
+ self._instrument_records[op.name] = TensorTracer.reason(
+ op_id, _REASON_USER_INCLUDED)
+ return False
+ if user_excluded:
+ self._instrument_records[op.name] = TensorTracer.reason(
+ op_id, _REASON_USER_EXCLUDED)
+ return True
+ if not self._inside_op_range(op_id):
+ self._instrument_records[op.name] = TensorTracer.reason(
+ op_id, _REASON_OUTSIDE_OP_RANGE)
+ return True
+ if TensorTracer.unsafe_op(op):
+ self._instrument_records[op.name] = TensorTracer.reason(
+ op_id, _REASON_UNSAFE_OP)
+ return True
+ if TensorTracer.device_mismatch(self._device_type, op):
+ self._instrument_records[op.name] = TensorTracer.reason(
+ op_id, _REASON_DEVICE_MISMATCH)
+ return True
+ if TensorTracer.less_interesting_op(op):
+ self._instrument_records[op.name] = TensorTracer.reason(
+ op_id, _REASON_LESS_INTERESTING_OP)
+ return True
+ return False
+
+ def _skip_tensor(self, op_id, out_tensor, user_included,
+ user_excluded):
+ """Returns True if we should not trace out_tensor."""
+
+ # Skips a tensor if the tensor has a non-numeric type.
+ # Note: we cannot use check_ops.is_numeric_tensor(out_tensor)
+ # because it also excludes tensors with dtypes, bool, and
+ # float32_ref, which we actually want to trace.
+ non_numeric_tensor_types = set([dtypes.variant, dtypes.resource,
+ dtypes.string])
+ if out_tensor.dtype in non_numeric_tensor_types:
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_NON_NUMERIC_TENSOR)
+ return True
+
+ if user_included:
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_USER_INCLUDED)
+ return False
+ if user_excluded:
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_USER_EXCLUDED)
+ return True
+ if not out_tensor.get_shape().is_fully_defined():
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_DYNAMIC_SHAPE)
+ return True
+ rank = len(out_tensor.shape)
+ if rank < 1:
+ # scalar
+ if TensorTracer.unsafe_scalar_trace(out_tensor.op):
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_UNSAFE_SCALAR)
+ return True
+ else:
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_SCALAR_GET_TRACED)
+ return False
+ else:
+ # tensor
+ self._instrument_records[out_tensor.name] = TensorTracer.reason(
+ op_id, _REASON_TENSOR_GET_TRACED)
+ return False
+
+ def _pre_tracing(self, graph):
+ """Work needs to be done prior to TPU or CPU tracing."""
+
+ operations = graph.get_operations()
+ (opname_idx_map, tensor_list, tensorname_idx_map) = (
+ TensorTracer._make_op_and_tensor_maps(operations))
+ self._write_config_section()
+ self._write_op_list_section(operations, tensorname_idx_map)
+ self._write_tensor_list_section(tensor_list, opname_idx_map)
+ # Does the topological sort before adding any nodes to the graph.
+ (succeed, sorted_or_cycle) = TensorTracer.topological_sort(graph)
+ return (operations, succeed, sorted_or_cycle)
+
+ def _post_tracing(self, succeed, sorted_or_cycle):
+ """Work needs to be done after TPU or CPU tracing."""
+
+ self._write_reason_section()
+ self._write_graph_section(succeed, sorted_or_cycle)
+ self._close_report_file()
+
+ def _get_checkpoints(self, graph):
+ """Returns the list of Ops that produce the tensors traced with API.
+
+ Args:
+ graph: the graph of Ops.
+
+ Returns:
+ A set of operation names which should be traced.
+ """
+
+ self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN,
+ _TENSOR_TRACER_CHECKPOINT))
+ checkpoint_operations = set()
+ tensor_tracer_variables = graph.get_collection(_TENSOR_TRACER_COLLECTION)
+ for (tensor, checkpoint_name) in tensor_tracer_variables:
+ self._write_report('%s %s\n'%(tensor.name, checkpoint_name))
+ checkpoint_operations.add(tensor.op.name)
+ self._write_report('%s %s\n'%(_MARKER_SECTION_END,
+ _TENSOR_TRACER_CHECKPOINT))
+ return checkpoint_operations
+
def trace_tpu(self, graph, result_tensor, num_replicas=None):
"""Traces the tensors generated by TPU Ops in a TF graph.
Args:
- graph: the graph of Ops.
+ graph: the graph of Ops executed on the TPU.
result_tensor: a result tensor of evaluating the graph.
num_replicas: number of replicas used on the TPU.
@@ -502,38 +880,22 @@
TensorTracer.check_device_type(self._device_type)
result_tensor_copy = self._add_replica_id_to_graph(num_replicas,
result_tensor)
- self._write_config_section()
+ (operations, succeed, sorted_or_cycle) = self._pre_tracing(graph)
tracing_ops = []
- operations = graph.get_operations()
- self._write_op_list_section(operations)
- # Does the topological sort before adding any nodes to the graph.
- (succeed, sorted_or_cycle) = TensorTracer.topological_sort(graph)
+ checkpoint_operations = self._get_checkpoints(graph)
+
for op_id, op in enumerate(operations):
- if not self._inside_op_range(op_id):
- self._instrument_records[op.name] = TensorTracer.reason(
- op_id, _RECORD_OUTSIDE_OP_RANGE)
+ if checkpoint_operations and op.name not in checkpoint_operations:
continue
- if not TensorTracer.should_trace(self._device_type, op):
- self._instrument_records[op.name] = TensorTracer.reason(
- op_id, _RECORD_SHOULD_NOT_TRACE)
- continue
- if not self._is_selected_op(op.name):
- self._instrument_records[op.name] = TensorTracer.reason(
- op_id, _RECORD_FILTERED_OUT)
+ user_included = self._is_user_included_op(op)
+ user_excluded = self._is_user_excluded_op(op)
+ if self._skip_op(op_id, op, user_included, user_excluded):
continue
for i in range(len(op.outputs)):
out_tensor = op.outputs[i]
- if not out_tensor.get_shape().is_fully_defined():
- self._instrument_records[out_tensor.name] = TensorTracer.reason(
- op_id, _RECORD_DYNAMIC_SHAPE)
- continue # cannot trace tensors with dynamic shape.
- rank = len(out_tensor.shape)
- if rank < 1:
- self._instrument_records[out_tensor.name] = TensorTracer.reason(
- op_id, _RECORD_SCALAR)
- continue # cannot trace scalar.
- self._instrument_records[out_tensor.name] = TensorTracer.reason(
- op_id, _RECORD_GET_TRACED)
+ if self._skip_tensor(op_id, out_tensor, user_included,
+ user_excluded):
+ continue
consumers = out_tensor.consumers()
trace_op = tpu.outside_compilation(
self._make_tensor_trace_fun(op.name, i), out_tensor)
@@ -546,8 +908,45 @@
# if there is no consumer, we will add the control dependence later
# when we add the control dependency to the output operations.
tracing_ops.append(trace_op)
-
- self._write_reason_section()
- self._write_graph_section(succeed, sorted_or_cycle)
-
+ self._post_tracing(succeed, sorted_or_cycle)
return (result_tensor_copy, tracing_ops)
+
+ def trace_cpu(self, graph):
+ """Traces the tensors generated by CPU Ops in a TF graph.
+
+ Args:
+ graph: the graph of Ops executed on the CPU.
+
+ Returns:
+ tracing_calls: a map from keys to trace calls.
+ A key is constructed from an Op's name.
+ A trace call consists of a function and a tensor (
+ the function will be invoked with the tensor).
+ """
+
+ self._device_type = _DEVICE_TYPE_CPU
+ TensorTracer.check_device_type(self._device_type)
+ self._num_replicas = 1
+ self._replica_id = 0
+ (operations, succeed, sorted_or_cycle) = self._pre_tracing(graph)
+ tracing_calls = {}
+ checkpoint_operations = self._get_checkpoints(graph)
+
+ for op_id, op in enumerate(operations):
+ if checkpoint_operations and op.name not in checkpoint_operations:
+ continue
+ user_included = self._is_user_included_op(op)
+ user_excluded = self._is_user_excluded_op(op)
+ if self._skip_op(op_id, op, user_included, user_excluded):
+ continue
+ for i in range(len(op.outputs)):
+ out_tensor = op.outputs[i]
+ if self._skip_tensor(op_id, out_tensor, user_included,
+ user_excluded):
+ continue
+ trace_fun = self._make_tensor_trace_fun(op.name, i)
+ trace_call = (trace_fun, [out_tensor])
+ trace_call_key = 'tensor_tracing_cpu-%s:%d'%(op.name, i)
+ tracing_calls[trace_call_key] = trace_call
+ self._post_tracing(succeed, sorted_or_cycle)
+ return tracing_calls
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py
index def57da..9266d81 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu.py
@@ -646,6 +646,10 @@
array_ops.identity(x, name="replicated_input_{}".format(i))
for i, x in enumerate(computation_inputs)
]
+ for i in computation_inputs:
+ # pylint: disable=protected-access
+ i.op._set_attr("_tpu_input_identity", attr_value_pb2.AttrValue(b=True))
+ # pylint: enable=protected-access
# If there is an infeed queue, adds the dequeued values to the
# computation's inputs.
@@ -726,7 +730,11 @@
new_output_tensors = []
for t in output_tensors:
with ops.device(t.device if t.device else core(0)):
- new_output_tensors.append(array_ops.identity(t))
+ o = array_ops.identity(t)
+ # pylint: disable=protected-access
+ o.op._set_attr("_tpu_output_identity", attr_value_pb2.AttrValue(b=True))
+ # pylint: enable=protected-access
+ new_output_tensors.append(o)
output_tensors = new_output_tensors
context.ExitResult(output_tensors)
finally:
@@ -777,6 +785,157 @@
]
+def split_compile_and_shard(computation,
+ inputs=None,
+ num_shards=1,
+ input_shard_axes=None,
+ outputs_from_all_shards=True,
+ output_shard_axes=None,
+ infeed_queue=None,
+ device_assignment=None,
+ name=None):
+ """Shards `computation` for parallel execution.
+
+ `inputs` must be a list of Tensors or None (equivalent to an empty list), each
+ of which has a corresponding split axis (from `input_shard_axes`). Each input
+ is split into `num_shards` pieces along the corresponding axis, and
+ computation is applied to each shard in parallel.
+
+ Tensors are broadcast to all shards if they are lexically captured by
+ `computation`. e.g.,
+
+ x = tf.constant(7)
+ def computation():
+ return x + 3
+ ... = shard(computation, ...)
+
+ TODO(phawkins): consider adding support for broadcasting Tensors passed
+ as inputs.
+
+ If `outputs_from_all_shards` is true, the outputs from all shards of
+ `computation` are concatenated back together along their `output_shards_axes`.
+ Otherwise, each output is taken from an arbitrary shard.
+
+ Inputs and outputs of the computation must be at least rank-1 Tensors.
+
+ Args:
+ computation: A Python function that builds a computation to apply to each
+ shard of the input.
+ inputs: A list of input tensors or None (equivalent to an empty list). Each
+ input tensor has a corresponding shard axes, given by `input_shard_axes`,
+ which must have size divisible by `num_shards`.
+ num_shards: The number of shards.
+ input_shard_axes: A list of dimensions along which to shard `inputs`, or
+ `None`. `None` means "shard all inputs along dimension 0". If not `None`,
+ there must be one dimension per input.
+ outputs_from_all_shards: Boolean or list of boolean. For each output, if
+ `True`, outputs from all shards are concatenated along the corresponding
+ `output_shard_axes` entry. Otherwise, each output is taken
+ from an arbitrary shard. If the argument is a boolean, the argument's
+ value is used for each output.
+ output_shard_axes: A list of dimensions along which to concatenate the
+ outputs of `computation`, or `None`. `None` means "concatenate all outputs
+ along dimension 0". If not `None`, there must be one dimension per output.
+ Ignored if `outputs_from_all_shards` is False.
+ infeed_queue: If not `None`, the `InfeedQueue` to use to augment the inputs
+ of `computation`.
+ device_assignment: If not `None`, a `DeviceAssignment` describing the
+ mapping between logical cores in the computation with physical cores in
+ the TPU topology. Uses a default device assignment if `None`. The
+ `DeviceAssignment` may be omitted if each shard of the computation uses
+ only one core, and there is either only one shard, or the number of shards
+ is equal to the number of cores in the TPU system.
+ name: (Deprecated) Does nothing.
+ Returns:
+ A tuple of (compile op, [output tensors]).
+ Raises:
+ ValueError: If num_shards <= 0
+ ValueError: If len(input_shard_axes) != len(inputs)
+ ValueError: If len(output_shard_axes) != len(outputs from `computation`)
+ """
+
+ if num_shards <= 0:
+ raise ValueError("num_shards must be a positive integer.")
+
+ inputs = [] if inputs is None else inputs
+ if not isinstance(inputs, list):
+ raise TypeError("tpu.shard()'s inputs must be a list of Tensors or None.")
+
+ # Converts inputs to Tensors.
+ inputs = [ops.convert_to_tensor(x) for x in inputs]
+
+ if input_shard_axes is None:
+ input_shard_axes = [0] * len(inputs)
+ if len(inputs) != len(input_shard_axes):
+ raise ValueError("Length of input_shard_axes must be equal to the number "
+ "of inputs.")
+
+ if inputs:
+ # Splits the `inputs` along the corresponding `input_shard_axes`, giving
+ # lists with layout [input][shard]
+ split_inputs = [
+ array_ops.split(x, num_shards, axis=axis)
+ for (axis, x) in zip(input_shard_axes, inputs)]
+
+ # Transposes the input lists to have layout [shard][input]
+ transposed_inputs = [list(i) for i in zip(*split_inputs)]
+ else:
+ transposed_inputs = [[]] * num_shards
+
+ compile_op, outputs = split_compile_and_replicate(
+ computation,
+ transposed_inputs,
+ infeed_queue=infeed_queue,
+ device_assignment=device_assignment,
+ name=name)
+
+ # There must be at least one shard since num_shards > 0.
+ # TODO(b/36647078) remove disable when pylint bug is fixed.
+ # pylint: disable=indexing-exception
+ if isinstance(outputs[0], ops.Operation):
+ # pylint: enable=indexing-exception
+ # There were no outputs from the computation and replicate returned a list
+ # of NoOps with control dependencies on the computation. Return the first
+ # one so it can be used as a control dependency or fetch node.
+ # TODO(b/36647078) remove disable when pylint bug is fixed.
+ # pylint: disable=indexing-exception
+ return compile_op, [outputs[0]]
+ # pylint: enable=indexing-exception
+
+ # TODO(b/36647078) remove disable when pylint bug is fixed.
+ # pylint: disable=indexing-exception
+ num_outputs = len(outputs[0])
+ # pylint: enable=indexing-exception
+
+ if output_shard_axes is None:
+ output_shard_axes = [0] * num_outputs
+ if num_outputs != len(output_shard_axes):
+ raise ValueError("Length of output_shard_axes must be equal to the number "
+ "of outputs.")
+
+ if isinstance(outputs_from_all_shards, bool):
+ outputs_from_all_shards = [outputs_from_all_shards] * num_outputs
+
+ if num_outputs != len(outputs_from_all_shards):
+ raise ValueError("Length of outputs_from_all_shards must be equal to the "
+ "number of outputs.")
+
+ results = []
+ for (axis, all_shards, x) in zip(output_shard_axes, outputs_from_all_shards,
+ zip(*outputs)):
+ if all_shards:
+ # Concatenate all of the outputs together (use stack for scalars).
+ shape = x[0].shape
+ is_scalar = shape is not None and (shape.ndims == 0)
+ results.append((array_ops.stack(list(x)) if is_scalar
+ else array_ops.concat(list(x), axis=axis)))
+ else:
+ # TODO(phawkins): use a smarter policy, e.g., round-robin across shards.
+ results.append(x[0])
+
+ return compile_op, results
+
+
def shard(computation,
inputs=None,
num_shards=1,
@@ -845,87 +1004,16 @@
ValueError: If len(input_shard_axes) != len(inputs)
ValueError: If len(output_shard_axes) != len(outputs from `computation`)
"""
-
- if num_shards <= 0:
- raise ValueError("num_shards must be a positive integer.")
-
- inputs = [] if inputs is None else inputs
- if not isinstance(inputs, list):
- raise TypeError("tpu.shard()'s inputs must be a list of Tensors or None.")
-
- # Converts inputs to Tensors.
- inputs = [ops.convert_to_tensor(x) for x in inputs]
-
- if input_shard_axes is None:
- input_shard_axes = [0] * len(inputs)
- if len(inputs) != len(input_shard_axes):
- raise ValueError("Length of input_shard_axes must be equal to the number "
- "of inputs.")
-
- if inputs:
- # Splits the `inputs` along the corresponding `input_shard_axes`, giving
- # lists with layout [input][shard]
- split_inputs = [
- array_ops.split(x, num_shards, axis=axis)
- for (axis, x) in zip(input_shard_axes, inputs)]
-
- # Transposes the input lists to have layout [shard][input]
- transposed_inputs = [list(i) for i in zip(*split_inputs)]
- else:
- transposed_inputs = [[]] * num_shards
-
- outputs = replicate(
+ return split_compile_and_shard(
computation,
- transposed_inputs,
+ inputs=inputs,
+ num_shards=num_shards,
+ input_shard_axes=input_shard_axes,
+ outputs_from_all_shards=outputs_from_all_shards,
+ output_shard_axes=output_shard_axes,
infeed_queue=infeed_queue,
device_assignment=device_assignment,
- name=name)
-
- # There must be at least one shard since num_shards > 0.
- # TODO(b/36647078) remove disable when pylint bug is fixed.
- # pylint: disable=indexing-exception
- if isinstance(outputs[0], ops.Operation):
- # pylint: enable=indexing-exception
- # There were no outputs from the computation and replicate returned a list
- # of NoOps with control dependencies on the computation. Return the first
- # one so it can be used as a control dependency or fetch node.
- # TODO(b/36647078) remove disable when pylint bug is fixed.
- # pylint: disable=indexing-exception
- return [outputs[0]]
- # pylint: enable=indexing-exception
-
- # TODO(b/36647078) remove disable when pylint bug is fixed.
- # pylint: disable=indexing-exception
- num_outputs = len(outputs[0])
- # pylint: enable=indexing-exception
-
- if output_shard_axes is None:
- output_shard_axes = [0] * num_outputs
- if num_outputs != len(output_shard_axes):
- raise ValueError("Length of output_shard_axes must be equal to the number "
- "of outputs.")
-
- if isinstance(outputs_from_all_shards, bool):
- outputs_from_all_shards = [outputs_from_all_shards] * num_outputs
-
- if num_outputs != len(outputs_from_all_shards):
- raise ValueError("Length of outputs_from_all_shards must be equal to the "
- "number of outputs.")
-
- results = []
- for (axis, all_shards, x) in zip(output_shard_axes, outputs_from_all_shards,
- zip(*outputs)):
- if all_shards:
- # Concatenate all of the outputs together (use stack for scalars).
- shape = x[0].shape
- is_scalar = shape is not None and (shape.ndims == 0)
- results.append((array_ops.stack(list(x)) if is_scalar
- else array_ops.concat(list(x), axis=axis)))
- else:
- # TODO(phawkins): use a smarter policy, e.g., round-robin across shards.
- results.append(x[0])
-
- return results
+ name=name)[1]
def batch_parallel(computation,
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index a9dc542..44a8f7c 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -31,6 +31,7 @@
from six.moves import queue as Queue # pylint: disable=redefined-builtin
from six.moves import xrange # pylint: disable=redefined-builtin
+from tensorflow.contrib.tpu.proto import compilation_result_pb2 as tpu_compilation_result
from tensorflow.contrib.tpu.python.tpu import tensor_tracer
from tensorflow.contrib.tpu.python.ops import tpu_ops
from tensorflow.contrib.tpu.python.tpu import error_handling
@@ -45,6 +46,7 @@
from tensorflow.core.framework import variable_pb2
from tensorflow.core.framework.summary_pb2 import Summary
from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session as tf_session
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest as data_nest
from tensorflow.python.estimator import estimator as estimator_lib
@@ -335,6 +337,16 @@
hooks = None
if self.host_call is not None:
hooks = [_OutfeedHostCallHook(host_call_ret['host_call'])]
+ if tensor_tracer.TensorTracer.is_enabled():
+ tt = tensor_tracer.TensorTracer()
+ tracing_calls = tt.trace_cpu(ops.get_default_graph())
+ tracing_call_ret = _OutfeedHostCall.create_cpu_hostcall(tracing_calls)
+ tracing_functions = tracing_call_ret.values()
+ if tracing_functions:
+ if hooks:
+ hooks.extend([_OutfeedHostCallHook(tracing_functions)])
+ else:
+ hooks = [_OutfeedHostCallHook(tracing_functions)]
hooks = tuple(hooks or [])
scaffold = self.scaffold_fn() if self.scaffold_fn else None
return model_fn_lib.EstimatorSpec(
@@ -411,13 +423,17 @@
ctx,
enqueue_ops,
dequeue_ops,
+ tpu_compile_op,
run_infeed_loop_on_coordinator=True,
- rendezvous=None):
+ rendezvous=None,
+ master=None,
+ session_config=None):
self._master_job = ctx.master_job
self._enqueue_ops = enqueue_ops
self._dequeue_ops = dequeue_ops
self._rendezvous = rendezvous
-
+ self._master = master
+ self._session_config = session_config
self._run_infeed_loop_on_coordinator = run_infeed_loop_on_coordinator
self._initial_infeed_sleep_secs = (
ctx.config.tpu_config.initial_infeed_sleep_secs)
@@ -425,15 +441,15 @@
self._feed_error = None
self._finished = False
self._should_initialize_tpu = True
+ self._tpu_compile_op = tpu_compile_op
def begin(self):
logging.info('TPU job name %s', self._master_job)
self._iterations_per_loop_var = _create_or_get_iterations_per_loop()
+ self._init_ops = []
if self._should_initialize_tpu:
- self._init_ops = [tpu.initialize_system(job=self._master_job)]
self._finalize_ops = [tpu.shutdown_system(job=self._master_job)]
else:
- self._init_ops = []
self._finalize_ops = []
summary_writer_init_ops = contrib_summary.summary_writer_initializer_op()
@@ -474,12 +490,31 @@
def _create_infeed_controller(self, name, target, args):
return _OpQueueContext(name=name, target=target, args=args)
+ def _assertCompilationSucceeded(self, result, coord):
+ proto = tpu_compilation_result.CompilationResultProto()
+ proto.ParseFromString(result)
+ if proto.status_error_message:
+ logging.error('Compilation failed: {}'.format(proto.status_error_message))
+ coord.request_stop()
+ else:
+ logging.info('Compilation succeeded')
+
def after_create_session(self, session, coord):
- logging.info('Init TPU system')
- start = time.time()
+ if self._should_initialize_tpu:
+ logging.info('Init TPU system')
+ start = time.time()
+ with ops.Graph().as_default():
+ with tf_session.Session(
+ self._master, config=self._session_config) as sess:
+ sess.run(tpu.initialize_system(job=self._master_job))
+ logging.info('Initialized TPU in %d seconds', time.time() - start)
+
session.run(self._init_ops,
options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000))
- logging.info('Initialized TPU in %d seconds', time.time() - start)
+
+ if os.environ.get('TPU_SPLIT_COMPILE_AND_EXECUTE', '') == '1':
+ logging.info('Compiling user program: this may take a while...')
+ self._assertCompilationSucceeded(session.run(self._tpu_compile_op), coord)
self._infeed_controller = self._create_infeed_controller(
name='InfeedController', target=self._run_infeed, args=(session,))
@@ -521,13 +556,17 @@
class TPUInfeedOutfeedSessionHookForPrediction(TPUInfeedOutfeedSessionHook):
- def __init__(self, ctx, enqueue_ops, dequeue_ops, rendezvous=None):
+ def __init__(self, ctx, enqueue_ops, dequeue_ops, tpu_compile_op,
+ rendezvous=None, master=None, session_config=None):
super(TPUInfeedOutfeedSessionHookForPrediction, self).__init__(
ctx,
enqueue_ops,
dequeue_ops,
+ tpu_compile_op=tpu_compile_op,
run_infeed_loop_on_coordinator=False,
- rendezvous=rendezvous)
+ rendezvous=rendezvous,
+ master=master,
+ session_config=session_config)
def _create_infeed_controller(self, name, target, args):
return _OpSignalOnceQueueContext(name=name, target=target, args=args)
@@ -2241,7 +2280,7 @@
(k, _export_output_to_tensors(v))
for k, v in six.iteritems(estimator_spec.export_outputs))
tensors = nest.flatten(tensors_dict)
- tpu_tensors = [t for t in tensors if _is_tpu_tensor(t)]
+ tpu_tensors = [t for t in tensors if t is not None]
# We cannot return anything other than `tpu_tensors` here so we capture
# the rest for later use.
@@ -2255,18 +2294,10 @@
# `tpu_tensors_on_cpu`.
new_tensors = []
for t in tensors:
- if _is_tpu_tensor(t):
- new_tensors.append(tpu_tensors_on_cpu.pop(0))
- elif t is None:
+ if t is None:
new_tensors.append(None)
else:
- # Only fetching `tpu_tensors_on_cpu` does not trigger
- # TPU computation and blocks, so we add the control dependency here.
- control_inputs = (
- tpu_tensors_on_cpu if _is_iterable(tpu_tensors_on_cpu) else
- (tpu_tensors_on_cpu,))
- with ops.control_dependencies(control_inputs):
- new_tensors.append(array_ops.identity(t))
+ new_tensors.append(tpu_tensors_on_cpu.pop(0))
# Reconstruct `tensors_dict`.
new_tensors_dict = nest.pack_sequence_as(tensors_dict, new_tensors)
@@ -2523,7 +2554,7 @@
graph.add_to_collection(_TPU_ENQUEUE_OPS, enqueue_op)
if mode == model_fn_lib.ModeKeys.TRAIN:
- loss, host_call, scaffold, training_hooks = (
+ compile_op, loss, host_call, scaffold, training_hooks = (
_train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn))
host_ops = host_call.create_tpu_hostcall()
if host_ops is None:
@@ -2558,9 +2589,12 @@
ctx,
enqueue_ops,
host_ops,
+ tpu_compile_op=compile_op,
run_infeed_loop_on_coordinator=(
run_infeed_loop_on_coordinator),
rendezvous=self._rendezvous[mode],
+ master=self._config.master,
+ session_config=self._session_config,
),
InstallSignalHandlerHook()
])
@@ -2613,8 +2647,8 @@
scaffold=scaffold)
if mode == model_fn_lib.ModeKeys.EVAL:
- total_loss, host_calls, scaffold, eval_hooks = _eval_on_tpu_system(
- ctx, model_fn_wrapper, dequeue_fn)
+ compile_op, total_loss, host_calls, scaffold, eval_hooks = (
+ _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn))
iterations_per_loop_var = _create_or_get_iterations_per_loop()
mean_loss = math_ops.div(
total_loss,
@@ -2661,10 +2695,13 @@
ctx,
enqueue_ops,
eval_update_ops + host_ops,
+ tpu_compile_op=compile_op,
run_infeed_loop_on_coordinator=(
run_infeed_loop_on_coordinator),
- rendezvous=self._rendezvous[mode]),
- ] + input_hooks
+ rendezvous=self._rendezvous[mode],
+ master=self._config.evaluation_master,
+ session_config=self._session_config,
+ )] + input_hooks
if eval_hooks:
hooks.extend(eval_hooks)
@@ -2679,7 +2716,7 @@
# Predict
assert mode == model_fn_lib.ModeKeys.PREDICT
- (dummy_predict_op, host_calls,
+ (compile_op, dummy_predict_op, host_calls,
scaffold, prediction_hooks) = _predict_on_tpu_system(
ctx, model_fn_wrapper, dequeue_fn)
with ops.control_dependencies([dummy_predict_op]):
@@ -2735,7 +2772,10 @@
hooks = [
_StoppingPredictHook(scalar_stopping_signal),
TPUInfeedOutfeedSessionHookForPrediction(
- ctx, enqueue_ops, host_ops, rendezvous=self._rendezvous[mode]),
+ ctx, enqueue_ops, host_ops, rendezvous=self._rendezvous[mode],
+ tpu_compile_op=compile_op,
+ master=self._config.master,
+ session_config=self._session_config),
] + input_hooks
if prediction_hooks:
@@ -2750,17 +2790,6 @@
return _model_fn
-def _is_tpu_tensor(tensor):
- if not isinstance(tensor, ops.Tensor):
- return False
- try:
- tensor.op.get_attr(tpu._OUTSIDE_COMPILATION_ATTR) # pylint: disable=protected-access
- except ValueError:
- return True
- else:
- return False
-
-
def _export_output_to_tensors(export_output):
"""Get a list of `Tensors` used in `export_output`.
@@ -2832,15 +2861,16 @@
return training_loop.repeat(iterations_per_loop_var, single_tpu_eval_step,
[_ZERO_LOSS])
- (loss,) = tpu.shard(
+ (compile_op, loss,) = tpu.split_compile_and_shard(
multi_tpu_eval_steps_on_single_shard,
inputs=[],
num_shards=ctx.num_replicas,
outputs_from_all_shards=False,
device_assignment=ctx.device_assignment)
+ loss = loss[0]
scaffold = _get_scaffold(captured_scaffold_fn)
- return loss, host_calls, scaffold, captured_eval_hooks.get()
+ return compile_op, loss, host_calls, scaffold, captured_eval_hooks.get()
def _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn):
@@ -2855,15 +2885,16 @@
return training_loop.repeat(iterations_per_loop_var, single_tpu_train_step,
[_INITIAL_LOSS])
- (loss,) = tpu.shard(
+ (compile_op, loss,) = tpu.split_compile_and_shard(
multi_tpu_train_steps_on_single_shard,
inputs=[],
num_shards=ctx.num_replicas,
outputs_from_all_shards=False,
device_assignment=ctx.device_assignment)
+ loss = loss[0]
scaffold = _get_scaffold(captured_scaffold_fn)
- return loss, host_call, scaffold, captured_training_hooks.get()
+ return compile_op, loss, host_call, scaffold, captured_training_hooks.get()
def _predict_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn):
@@ -2883,15 +2914,17 @@
cond, single_tpu_predict_step, inputs=inputs, name=b'loop')
return outputs
- (dummy_predict_op,) = tpu.shard(
+ (compile_op, dummy_predict_op,) = tpu.split_compile_and_shard(
multi_tpu_predict_steps_on_single_shard,
inputs=[],
num_shards=ctx.num_replicas,
outputs_from_all_shards=False,
device_assignment=ctx.device_assignment)
+ dummy_predict_op = dummy_predict_op[0]
scaffold = _get_scaffold(captured_scaffold_fn)
- return dummy_predict_op, host_calls, scaffold, captured_predict_hooks.get()
+ return (compile_op, dummy_predict_op, host_calls, scaffold,
+ captured_predict_hooks.get())
def _wrap_computation_in_while_loop(device, op_fn):
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py
index 5523555..e3ea983 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py
@@ -21,8 +21,8 @@
import numpy as np
from tensorflow.contrib.tpu.python.tpu import tpu_estimator
-from tensorflow.python import data as dataset_lib
from tensorflow.python.client import session
+from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.platform import test
@@ -34,10 +34,10 @@
def input_fn(params):
batch_size = params['batch_size']
- da1 = dataset_lib.Dataset.from_tensor_slices(a)
- da2 = dataset_lib.Dataset.from_tensor_slices(b)
+ da1 = dataset_ops.Dataset.from_tensor_slices(a)
+ da2 = dataset_ops.Dataset.from_tensor_slices(b)
- dataset = dataset_lib.Dataset.zip((da1, da2))
+ dataset = dataset_ops.Dataset.zip((da1, da2))
dataset = dataset.map(lambda fa, fb: {'a': fa, 'b': fb})
dataset = dataset.batch(batch_size)
return dataset
@@ -50,10 +50,10 @@
def input_fn(params):
batch_size = params['batch_size']
- da1 = dataset_lib.Dataset.from_tensor_slices(a)
- da2 = dataset_lib.Dataset.from_tensor_slices(b)
+ da1 = dataset_ops.Dataset.from_tensor_slices(a)
+ da2 = dataset_ops.Dataset.from_tensor_slices(b)
- dataset = dataset_lib.Dataset.zip((da1, da2))
+ dataset = dataset_ops.Dataset.zip((da1, da2))
dataset = dataset.map(lambda fa, fb: ({'a': fa}, fb))
dataset = dataset.batch(batch_size)
return dataset
@@ -71,7 +71,7 @@
with ops.Graph().as_default():
dataset = input_fn(params)
- features = dataset_lib.make_one_shot_iterator(dataset).get_next()
+ features = dataset_ops.make_one_shot_iterator(dataset).get_next()
# With tf.data.Dataset.batch, the batch is None, i.e., dynamic shape.
self.assertIsNone(features['a'].shape.as_list()[0])
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py
index ec682e5..d66ecfc 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py
@@ -52,6 +52,7 @@
devices = []
device_dict = collections.defaultdict(list)
+ # TODO(b/120564445): Replace with standard library for retries.
retry_count = 1
while True:
logging.info('Querying Tensorflow master (%s) for TPU system metadata.',
diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py
index 3beb7bf..bcc1776 100644
--- a/tensorflow/contrib/training/python/training/hparam.py
+++ b/tensorflow/contrib/training/python/training/hparam.py
@@ -187,7 +187,7 @@
return param_type(value)
-def parse_values(values, type_map):
+def parse_values(values, type_map, ignore_unknown=False):
"""Parses hyperparameter values from a string into a python map.
`values` is a string containing comma-separated `name=value` pairs.
@@ -233,6 +233,9 @@
type T if either V has type T, or V is a list of elements of type T.
Hence, for a multidimensional parameter 'x' taking float values,
'x=[0.1,0.2]' will parse successfully if type_map['x'] = float.
+ ignore_unknown: Bool. Whether values that are missing a type in type_map
+ should be ignored. If set to True, a ValueError will not be raised for
+ unknown hyperparameter type.
Returns:
A python map mapping each name to either:
@@ -260,6 +263,8 @@
m_dict = m.groupdict()
name = m_dict['name']
if name not in type_map:
+ if ignore_unknown:
+ continue
raise ValueError('Unknown hyperparameter type for %s' % name)
type_ = type_map[name]
diff --git a/tensorflow/contrib/training/python/training/hparam_test.py b/tensorflow/contrib/training/python/training/hparam_test.py
index 660c97f..a990e04 100644
--- a/tensorflow/contrib/training/python/training/hparam_test.py
+++ b/tensorflow/contrib/training/python/training/hparam_test.py
@@ -216,6 +216,14 @@
self.assertTrue(isinstance(parse_dict['arr'], dict))
self.assertDictEqual(parse_dict['arr'], {1: 10})
+ def testParseValuesWithIndexAssigment1_IgnoreUnknown(self):
+ """Assignment to an index position."""
+ parse_dict = hparam.parse_values(
+ 'arr[1]=10,b=5', {'arr': int}, ignore_unknown=True)
+ self.assertEqual(len(parse_dict), 1)
+ self.assertTrue(isinstance(parse_dict['arr'], dict))
+ self.assertDictEqual(parse_dict['arr'], {1: 10})
+
def testParseValuesWithIndexAssigment2(self):
"""Assignment to multiple index positions."""
parse_dict = hparam.parse_values('arr[0]=10,arr[5]=20', {'arr': int})
@@ -223,6 +231,14 @@
self.assertTrue(isinstance(parse_dict['arr'], dict))
self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20})
+ def testParseValuesWithIndexAssigment2_IgnoreUnknown(self):
+ """Assignment to multiple index positions."""
+ parse_dict = hparam.parse_values(
+ 'arr[0]=10,arr[5]=20,foo=bar', {'arr': int}, ignore_unknown=True)
+ self.assertEqual(len(parse_dict), 1)
+ self.assertTrue(isinstance(parse_dict['arr'], dict))
+ self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20})
+
def testParseValuesWithIndexAssigment3(self):
"""Assignment to index positions in multiple names."""
parse_dict = hparam.parse_values('arr[0]=10,arr[1]=20,L[5]=100,L[10]=200',
@@ -234,6 +250,17 @@
self.assertTrue(isinstance(parse_dict['L'], dict))
self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200})
+ def testParseValuesWithIndexAssigment3_IgnoreUnknown(self):
+ """Assignment to index positions in multiple names."""
+ parse_dict = hparam.parse_values(
+ 'arr[0]=10,C=5,arr[1]=20,B[0]=kkk,L[5]=100,L[10]=200',
+ {'arr': int, 'L': int}, ignore_unknown=True)
+ self.assertEqual(len(parse_dict), 2)
+ self.assertTrue(isinstance(parse_dict['arr'], dict))
+ self.assertDictEqual(parse_dict['arr'], {0: 10, 1: 20})
+ self.assertTrue(isinstance(parse_dict['L'], dict))
+ self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200})
+
def testParseValuesWithIndexAssigment4(self):
"""Assignment of index positions and scalars."""
parse_dict = hparam.parse_values('x=10,arr[1]=20,y=30',
@@ -246,6 +273,17 @@
self.assertEqual(parse_dict['x'], 10)
self.assertEqual(parse_dict['y'], 30)
+ def testParseValuesWithIndexAssigment4_IgnoreUnknown(self):
+ """Assignment of index positions and scalars."""
+ parse_dict = hparam.parse_values(
+ 'x=10,foo[0]=bar,arr[1]=20,zzz=78,y=30',
+ {'x': int, 'y': int, 'arr': int}, ignore_unknown=True)
+ self.assertEqual(len(parse_dict), 3)
+ self.assertTrue(isinstance(parse_dict['arr'], dict))
+ self.assertDictEqual(parse_dict['arr'], {1: 20})
+ self.assertEqual(parse_dict['x'], 10)
+ self.assertEqual(parse_dict['y'], 30)
+
def testParseValuesWithIndexAssigment5(self):
"""Different variable types."""
parse_dict = hparam.parse_values('a[0]=5,b[1]=true,c[2]=abc,d[3]=3.14', {
@@ -264,24 +302,55 @@
self.assertTrue(isinstance(parse_dict['d'], dict))
self.assertDictEqual(parse_dict['d'], {3: 3.14})
+ def testParseValuesWithIndexAssigment5_IgnoreUnknown(self):
+ """Different variable types."""
+ parse_dict = hparam.parse_values(
+ 'a[0]=5,cc=4,b[1]=true,c[2]=abc,mm=2,d[3]=3.14',
+ {'a': int, 'b': bool, 'c': str, 'd': float},
+ ignore_unknown=True)
+ self.assertEqual(set(parse_dict.keys()), {'a', 'b', 'c', 'd'})
+ self.assertTrue(isinstance(parse_dict['a'], dict))
+ self.assertDictEqual(parse_dict['a'], {0: 5})
+ self.assertTrue(isinstance(parse_dict['b'], dict))
+ self.assertDictEqual(parse_dict['b'], {1: True})
+ self.assertTrue(isinstance(parse_dict['c'], dict))
+ self.assertDictEqual(parse_dict['c'], {2: 'abc'})
+ self.assertTrue(isinstance(parse_dict['d'], dict))
+ self.assertDictEqual(parse_dict['d'], {3: 3.14})
+
def testParseValuesWithBadIndexAssigment1(self):
"""Reject assignment of list to variable type."""
with self.assertRaisesRegexp(ValueError,
r'Assignment of a list to a list index.'):
hparam.parse_values('arr[1]=[1,2,3]', {'arr': int})
+ def testParseValuesWithBadIndexAssigment1_IgnoreUnknown(self):
+ """Reject assignment of list to variable type."""
+ with self.assertRaisesRegexp(ValueError,
+ r'Assignment of a list to a list index.'):
+ hparam.parse_values(
+ 'arr[1]=[1,2,3],c=8', {'arr': int}, ignore_unknown=True)
+
def testParseValuesWithBadIndexAssigment2(self):
"""Reject if type missing."""
with self.assertRaisesRegexp(ValueError,
r'Unknown hyperparameter type for arr'):
hparam.parse_values('arr[1]=5', {})
+ def testParseValuesWithBadIndexAssigment2_IgnoreUnknown(self):
+ """Ignore missing type."""
+ hparam.parse_values('arr[1]=5', {}, ignore_unknown=True)
+
def testParseValuesWithBadIndexAssigment3(self):
"""Reject type of the form name[index]."""
with self.assertRaisesRegexp(ValueError,
'Unknown hyperparameter type for arr'):
hparam.parse_values('arr[1]=1', {'arr[1]': int})
+ def testParseValuesWithBadIndexAssigment3_IgnoreUnknown(self):
+ """Ignore type of the form name[index]."""
+ hparam.parse_values('arr[1]=1', {'arr[1]': int}, ignore_unknown=True)
+
def testWithReusedVariables(self):
with self.assertRaisesRegexp(ValueError,
'Multiple assignments to variable \'x\''):
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 575edfe..2760050 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -49,7 +49,7 @@
# filegroup ":android_proto_srcs" - Protos
# filegroup ":android_srcs" - Core sources
# cc_library ":android_tensorflow_lib" - Native library
-# cc_library ":android_tensorflow_lib_selective_registration" - Native library
+# cc_library ":android_tensorflow_lib_lite" - Native library, without ops,
# supporting SELECTIVE_REGISTRATION feature.
# portable_proto_library ":android_proto_lib" (Google-internal)
#
@@ -446,15 +446,19 @@
)
cc_library(
- name = "logger",
- srcs = tf_platform_srcs(["logger.cc"]),
- hdrs = ["platform/logger.h"] + tf_platform_hdrs(["logger.h"]),
+ name = "logger_interface",
+ hdrs = ["platform/logger.h"],
copts = tf_copts(),
visibility = ["//visibility:public"],
- deps = [
- ":lib",
- ":lib_internal",
- ] + tf_additional_logger_deps(),
+ deps = [":platform_protobuf"],
+)
+
+cc_library(
+ name = "logger",
+ srcs = tf_platform_srcs(["logger.cc"]),
+ copts = tf_copts(),
+ visibility = ["//visibility:public"],
+ deps = [":logger_interface"] + tf_additional_logger_deps(),
)
filegroup(
@@ -492,7 +496,10 @@
":platform_env_internal_hdrs",
],
copts = tf_copts(),
- visibility = ["//tensorflow/core:__subpackages__"],
+ visibility = [
+ "//tensorflow/c:__subpackages__",
+ "//tensorflow/core:__subpackages__",
+ ],
deps = [
":error_codes_proto_cc",
":lib",
@@ -1616,7 +1623,6 @@
"util/reporter.*",
"platform/**/cuda_libdevice_path.*",
"platform/**/logger.cc",
- "platform/**/logger.h",
"platform/default/test_benchmark.*",
"platform/cuda.h",
"platform/google/**/*",
@@ -1679,6 +1685,9 @@
# operators, use :android_tensorflow_lib if you want full operator
# support.
#
+# If you just need TensorFlow types, e.g. Tensors, use
+# :android_tensorflow_lib_lite_no_runtime.
+#
# Compiles to a trivial library on non-Android to prevent irrelevant
# build errors. If not building this as part of an android_binary,
# a command such as the following must be used:
@@ -1689,7 +1698,33 @@
cc_library(
name = "android_tensorflow_lib_lite",
srcs = if_android(["//tensorflow/core:android_srcs"]),
- copts = tf_copts(android_optimization_level_override = None),
+ copts = tf_copts(android_optimization_level_override = None) + [
+ "-DSUPPORT_SELECTIVE_REGISTRATION",
+ ],
+ linkopts = ["-lz"],
+ tags = [
+ "manual",
+ "notap",
+ ],
+ visibility = ["//visibility:public"],
+ deps = [
+ ":mobile_additional_lib_deps",
+ ":protos_all_cc_impl",
+ ":stats_calculator_portable",
+ "//third_party/eigen3",
+ "@double_conversion//:double-conversion",
+ "@nsync//:nsync_cpp",
+ "@protobuf_archive//:protobuf",
+ ],
+ alwayslink = 1,
+)
+
+cc_library(
+ name = "android_tensorflow_lib_lite_nortti",
+ srcs = if_android(["//tensorflow/core:android_srcs"]),
+ copts = tf_copts(android_optimization_level_override = None) + [
+ "-DSUPPORT_SELECTIVE_REGISTRATION",
+ ] + tf_opts_nortti_if_android(),
linkopts = ["-lz"],
tags = [
"manual",
@@ -1797,52 +1832,6 @@
alwayslink = 1,
)
-# Android library for use with the SELECTIVE_REGISTRATION feature.
-# Does not contain operators. In contrast to android_tensorflow_lib_lite,
-# this links in framework support for all types, relying on selective
-# registration of ops to prune code size.
-cc_library(
- name = "android_tensorflow_lib_selective_registration",
- srcs = if_android(["//tensorflow/core:android_srcs_only_runtime"]),
- copts = tf_copts(android_optimization_level_override = None) + [
- "-DSUPPORT_SELECTIVE_REGISTRATION",
- ],
- linkopts = if_android(["-lz"]),
- tags = [
- "manual",
- "notap",
- ],
- visibility = ["//visibility:public"],
- deps = [
- ":protos_all_cc_impl",
- "@com_google_absl//absl/container:flat_hash_set",
- "@protobuf_archive//:protobuf",
- ],
- alwayslink = 1,
-)
-
-# Android library for use with the SELECTIVE_REGISTRATION feature with
-# no proto_rtti.
-cc_library(
- name = "android_tensorflow_lib_selective_registration_nortti",
- srcs = if_android(["//tensorflow/core:android_srcs_only_runtime"]),
- copts = tf_copts(android_optimization_level_override = None) + tf_opts_nortti_if_android() + [
- "-DSUPPORT_SELECTIVE_REGISTRATION",
- ],
- linkopts = if_android(["-lz"]),
- tags = [
- "manual",
- "notap",
- ],
- visibility = ["//visibility:public"],
- deps = [
- ":protos_all_cc_impl",
- "@com_google_absl//absl/container:flat_hash_set",
- "@protobuf_archive//:protobuf",
- ],
- alwayslink = 1,
-)
-
filegroup(
name = "android_op_registrations_and_gradients",
srcs = glob(
@@ -4052,20 +4041,6 @@
],
)
-tf_cc_test_gpu(
- name = "cuda_libdevice_path_test",
- size = "small",
- srcs = ["platform/cuda_libdevice_path_test.cc"],
- linkstatic = tf_kernel_tests_linkstatic(),
- tags = tf_cuda_tests_tags(),
- deps = [
- ":cuda_libdevice_path",
- ":lib",
- ":test",
- ":test_main",
- ],
-)
-
tf_cuda_only_cc_test(
name = "util_cuda_kernel_helper_test",
srcs = [
@@ -4921,7 +4896,7 @@
cc_library(
name = "cuda_libdevice_path",
- srcs = ["platform/cuda_libdevice_path.cc"] + tf_additional_libdevice_srcs(),
+ srcs = tf_additional_libdevice_srcs(),
hdrs = ["platform/cuda_libdevice_path.h"],
copts = tf_copts(),
data = tf_additional_libdevice_data(),
diff --git a/tensorflow/core/api_def/api_test.cc b/tensorflow/core/api_def/api_test.cc
index d38a842..7405e2a 100644
--- a/tensorflow/core/api_def/api_test.cc
+++ b/tensorflow/core/api_def/api_test.cc
@@ -35,7 +35,6 @@
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
-#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
namespace {
diff --git a/tensorflow/core/api_def/base_api/api_def_Lu.pbtxt b/tensorflow/core/api_def/base_api/api_def_Lu.pbtxt
new file mode 100644
index 0000000..35dbee8
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_Lu.pbtxt
@@ -0,0 +1,51 @@
+op {
+ graph_op_name: "Lu"
+ in_arg {
+ name: "input"
+ description: <<END
+A tensor of shape `[..., M, M]` whose inner-most 2 dimensions form matrices of
+size `[M, M]`.
+END
+ }
+ out_arg {
+ name: "lu"
+ description: <<END
+A tensor of shape `[..., M, M]` whose strictly lower triangular part denotes the
+lower triangular factor `L` with unit diagonal, and whose upper triangular part
+denotes the upper triangular factor `U`.
+END
+ }
+ out_arg {
+ name: "p"
+ description: <<END
+Permutation of the rows encoded as a list of indices in `0..M-1`. Shape is
+`[..., M]`.
+@compatibility(scipy)
+Similar to `scipy.linalg.lu`, except the triangular factors `L` and `U` are
+packed into a single tensor, the permutation is applied to `input` instead of
+the right hand side and the permutation `P` is returned as a list of indices
+instead of a permutation matrix.
+@end_compatibility
+END
+ }
+ summary: "Computes the LU decomposition of one or more square matrices."
+ description: <<END
+The input is a tensor of shape `[..., M, M]` whose inner-most 2 dimensions
+form square matrices.
+
+The input has to be invertible.
+
+The output consists of two tensors LU and P containing the LU decomposition
+of all input submatrices `[..., :, :]`. LU encodes the lower triangular and
+upper triangular factors.
+
+For each input submatrix of shape `[M, M]`, L is a lower triangular matrix of
+shape `[M, M]` with unit diagonal whose entries correspond to the strictly lower
+triangular part of LU. U is a upper triangular matrix of shape `[M, M]` whose
+entries correspond to the upper triangular part, including the diagonal, of LU.
+
+P represents a permutation matrix encoded as a list of indices each between `0`
+and `M-1`, inclusive. If P_mat denotes the permutation matrix corresponding to
+P, then the L, U and P satisfies P_mat * input = L * U.
+END
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_OneHot.pbtxt b/tensorflow/core/api_def/base_api/api_def_OneHot.pbtxt
index 807b8ae..b325df1 100644
--- a/tensorflow/core/api_def/base_api/api_def_OneHot.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_OneHot.pbtxt
@@ -66,7 +66,6 @@
=========
Suppose that
-
```
indices = [0, 2, -1, 1]
depth = 3
@@ -76,16 +75,15 @@
```
Then output is `[4 x 3]`:
-
- ```output =
- [5.0 0.0 0.0] // one_hot(0)
- [0.0 0.0 5.0] // one_hot(2)
- [0.0 0.0 0.0] // one_hot(-1)
- [0.0 5.0 0.0] // one_hot(1)
- ```
+```
+output =
+ [5.0 0.0 0.0] // one_hot(0)
+ [0.0 0.0 5.0] // one_hot(2)
+ [0.0 0.0 0.0] // one_hot(-1)
+ [0.0 5.0 0.0] // one_hot(1)
+```
Suppose that
-
```
indices = [0, 2, -1, 1]
depth = 3
@@ -95,19 +93,19 @@
```
Then output is `[3 x 4]`:
+```
+output =
+ [0.0 3.0 3.0 3.0]
+ [3.0 3.0 3.0 0.0]
+ [3.0 3.0 3.0 3.0]
+ [3.0 0.0 3.0 3.0]
+// ^ one_hot(0)
+// ^ one_hot(2)
+// ^ one_hot(-1)
+// ^ one_hot(1)
+```
- ```output =
- [0.0 3.0 3.0 3.0]
- [3.0 3.0 3.0 0.0]
- [3.0 3.0 3.0 3.0]
- [3.0 0.0 3.0 3.0]
- // ^ one_hot(0)
- // ^ one_hot(2)
- // ^ one_hot(-1)
- // ^ one_hot(1)
- ```
Suppose that
-
```
indices = [[0, 2], [1, -1]]
depth = 3
@@ -117,14 +115,15 @@
```
Then output is `[2 x 2 x 3]`:
-
- ```output =
- [
- [1.0, 0.0, 0.0] // one_hot(0)
- [0.0, 0.0, 1.0] // one_hot(2)
- ][
- [0.0, 1.0, 0.0] // one_hot(1)
- [0.0, 0.0, 0.0] // one_hot(-1)
- ]```
+```
+output =
+ [
+ [1.0, 0.0, 0.0] // one_hot(0)
+ [0.0, 0.0, 1.0] // one_hot(2)
+ ][
+ [0.0, 1.0, 0.0] // one_hot(1)
+ [0.0, 0.0, 0.0] // one_hot(-1)
+ ]
+```
END
}
diff --git a/tensorflow/core/api_def/base_api/api_def_RaggedGather.pbtxt b/tensorflow/core/api_def/base_api/api_def_RaggedGather.pbtxt
index 240c987..9c40332 100644
--- a/tensorflow/core/api_def/base_api/api_def_RaggedGather.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_RaggedGather.pbtxt
@@ -11,8 +11,8 @@
in_arg {
name: "params_dense_values"
description: <<END
-The `inner_values` for the `params` RaggedTensor. There was a terminology change
-at the python level from dense_values to inner_values, so dense_values is the
+The `flat_values` for the `params` RaggedTensor. There was a terminology change
+at the python level from dense_values to flat_values, so dense_values is the
deprecated name.
END
}
@@ -32,7 +32,7 @@
}
out_arg {
name: "output_dense_values"
- description: "The `inner_values` for the returned RaggedTensor."
+ description: "The `flat_values` for the returned RaggedTensor."
}
attr {
name: "PARAMS_RAGGED_RANK"
diff --git a/tensorflow/core/api_def/base_api/api_def_RaggedRange.pbtxt b/tensorflow/core/api_def/base_api/api_def_RaggedRange.pbtxt
index 927e839..4a9b2af 100644
--- a/tensorflow/core/api_def/base_api/api_def_RaggedRange.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_RaggedRange.pbtxt
@@ -19,7 +19,7 @@
}
out_arg{
name: "rt_dense_values"
- description: "The `inner_values` for the returned `RaggedTensor`."
+ description: "The `flat_values` for the returned `RaggedTensor`."
}
summary: <<END
Returns a `RaggedTensor` containing the specified sequences of numbers.
diff --git a/tensorflow/core/api_def/base_api/api_def_RaggedTensorToSparse.pbtxt b/tensorflow/core/api_def/base_api/api_def_RaggedTensorToSparse.pbtxt
index 8c73ea6..958c711 100644
--- a/tensorflow/core/api_def/base_api/api_def_RaggedTensorToSparse.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_RaggedTensorToSparse.pbtxt
@@ -7,7 +7,7 @@
}
in_arg {
name: "rt_dense_values"
- description: "The `inner_values` for the `RaggedTensor`."
+ description: "The `flat_values` for the `RaggedTensor`."
}
out_arg {
name: "sparse_indices"
diff --git a/tensorflow/core/api_def/base_api/api_def_UnicodeDecode.pbtxt b/tensorflow/core/api_def/base_api/api_def_UnicodeDecode.pbtxt
new file mode 100644
index 0000000..9b3f690
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_UnicodeDecode.pbtxt
@@ -0,0 +1,76 @@
+op {
+ graph_op_name: "UnicodeDecode"
+ in_arg {
+ name: "input"
+ description: <<END
+The text to be decoded. Can have any shape. Note that the output is flattened
+to a vector of char values.
+END
+ }
+ out_arg {
+ name: "row_splits"
+ description: <<END
+A 1D int32 tensor containing the row splits.
+END
+ }
+ out_arg {
+ name: "char_values"
+ description: <<END
+A 1D int32 Tensor containing the decoded codepoints.
+END
+ }
+ attr {
+ name: "input_encoding"
+ description: <<END
+Text encoding of the input strings. This is any of the encodings supported
+by ICU ucnv algorithmic converters. Examples: `"UTF-16", "US ASCII", "UTF-8"`.
+END
+ }
+ attr {
+ name: "errors"
+ description: <<END
+Error handling policy when there is invalid formatting found in the input.
+The value of 'strict' will cause the operation to produce a InvalidArgument
+error on any invalid input formatting. A value of 'replace' (the default) will
+cause the operation to replace any invalid formatting in the input with the
+`replacement_char` codepoint. A value of 'ignore' will cause the operation to
+skip any invalid formatting in the input and produce no corresponding output
+character.
+END
+ }
+ attr {
+ name: "replacement_char"
+ description: <<END
+The replacement character codepoint to be used in place of any invalid
+formatting in the input when `errors='replace'`. Any valid unicode codepoint may
+be used. The default value is the default unicode replacement character is
+0xFFFD or U+65533.)
+END
+ }
+ attr {
+ name: "replace_control_characters"
+ description: <<END
+Whether to replace the C0 control characters (00-1F) with the
+`replacement_char`. Default is false.
+END
+ }
+ summary: <<END
+Decodes each string in `input` into a sequence of Unicode code points.
+END
+ description: <<END
+The character codepoints for all strings are returned using a single vector
+`char_values`, with strings expanded to characters in row-major order.
+
+The `row_splits` tensor indicates where the codepoints for
+each input string begin and end within the `char_values` tensor.
+In particular, the values for the `i`th
+string (in row-major order) are stored in the slice
+`[row_splits[i]:row_splits[i+1]]`. Thus:
+
+* `char_values[row_splits[i]+j]` is the Unicode codepoint for the `j`th
+ character in the `i`th string (in row-major order).
+* `row_splits[i+1] - row_splits[i]` is the number of characters in the `i`th
+ string (in row-major order).
+END
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_UnwrapDatasetVariant.pbtxt b/tensorflow/core/api_def/base_api/api_def_UnwrapDatasetVariant.pbtxt
new file mode 100644
index 0000000..7b3f88a
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_UnwrapDatasetVariant.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "UnwrapDatasetVariant"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_WrapDatasetVariant.pbtxt b/tensorflow/core/api_def/base_api/api_def_WrapDatasetVariant.pbtxt
new file mode 100644
index 0000000..40f5c7a
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_WrapDatasetVariant.pbtxt
@@ -0,0 +1,4 @@
+op {
+ graph_op_name: "WrapDatasetVariant"
+ visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_FloorDiv.pbtxt b/tensorflow/core/api_def/python_api/api_def_FloorDiv.pbtxt
index 26598ab..efd42b8 100644
--- a/tensorflow/core/api_def/python_api/api_def_FloorDiv.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_FloorDiv.pbtxt
@@ -1,4 +1,6 @@
op {
graph_op_name: "FloorDiv"
- visibility: HIDDEN
+ endpoint {
+ name: "floor_div"
+ }
}
diff --git a/tensorflow/core/api_def/python_api/api_def_FloorMod.pbtxt b/tensorflow/core/api_def/python_api/api_def_FloorMod.pbtxt
index ef562e9..e5db6d4 100644
--- a/tensorflow/core/api_def/python_api/api_def_FloorMod.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_FloorMod.pbtxt
@@ -1,4 +1,9 @@
op {
graph_op_name: "FloorMod"
- visibility: HIDDEN
+ endpoint {
+ name: "floormod"
+ }
+ endpoint {
+ name: "mod"
+ }
}
diff --git a/tensorflow/core/api_def/python_api/api_def_Lu.pbtxt b/tensorflow/core/api_def/python_api/api_def_Lu.pbtxt
new file mode 100644
index 0000000..c0b6b53
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_Lu.pbtxt
@@ -0,0 +1,6 @@
+op {
+ graph_op_name: "Lu"
+ endpoint {
+ name: "linalg.lu"
+ }
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_RealDiv.pbtxt b/tensorflow/core/api_def/python_api/api_def_RealDiv.pbtxt
index bd87eef..f9e01eb 100644
--- a/tensorflow/core/api_def/python_api/api_def_RealDiv.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_RealDiv.pbtxt
@@ -1,4 +1,6 @@
op {
graph_op_name: "RealDiv"
- visibility: HIDDEN
+ endpoint {
+ name: "realdiv"
+ }
}
diff --git a/tensorflow/core/api_def/python_api/api_def_TruncateDiv.pbtxt b/tensorflow/core/api_def/python_api/api_def_TruncateDiv.pbtxt
index 2a547f7..8e46c5e 100644
--- a/tensorflow/core/api_def/python_api/api_def_TruncateDiv.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_TruncateDiv.pbtxt
@@ -1,4 +1,6 @@
op {
graph_op_name: "TruncateDiv"
- visibility: HIDDEN
+ endpoint {
+ name: "truncatediv"
+ }
}
diff --git a/tensorflow/core/api_def/python_api/api_def_TruncateMod.pbtxt b/tensorflow/core/api_def/python_api/api_def_TruncateMod.pbtxt
index 0731e88..97fb816 100644
--- a/tensorflow/core/api_def/python_api/api_def_TruncateMod.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_TruncateMod.pbtxt
@@ -1,4 +1,6 @@
op {
graph_op_name: "TruncateMod"
- visibility: HIDDEN
+ endpoint {
+ name: "truncatemod"
+ }
}
diff --git a/tensorflow/core/common_runtime/eager/context.cc b/tensorflow/core/common_runtime/eager/context.cc
index 583ae64..1727c04 100644
--- a/tensorflow/core/common_runtime/eager/context.cc
+++ b/tensorflow/core/common_runtime/eager/context.cc
@@ -15,6 +15,9 @@
#include "tensorflow/core/common_runtime/eager/context.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/device_resolver_local.h"
#include "tensorflow/core/common_runtime/device_set.h"
#include "tensorflow/core/common_runtime/process_util.h"
#include "tensorflow/core/framework/resource_mgr.h"
@@ -71,6 +74,13 @@
runner_ = [this](std::function<void()> closure) {
this->thread_pool_->Schedule(std::move(closure));
};
+
+ std::unique_ptr<DeviceResolverInterface> drl(
+ new DeviceResolverLocal(local_device_mgr()));
+ std::unique_ptr<ParamResolverInterface> cprl(new CollectiveParamResolverLocal(
+ local_device_mgr(), drl.get(), "/job:localhost/replica:0/task:0"));
+ collective_executor_mgr_.reset(new CollectiveExecutorMgr(
+ opts.config, local_device_mgr(), std::move(drl), std::move(cprl)));
}
void EagerContext::InitDeviceMapAndAsync() {
diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h
index 51109f8..cdef947 100644
--- a/tensorflow/core/common_runtime/eager/context.h
+++ b/tensorflow/core/common_runtime/eager/context.h
@@ -33,6 +33,7 @@
#include "tensorflow/core/distributed_runtime/eager/eager_client.h"
#include "tensorflow/core/distributed_runtime/server_lib.h"
#endif
+#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/framework/log_memory.h"
#include "tensorflow/core/framework/rendezvous.h"
#include "tensorflow/core/lib/core/stringpiece.h"
@@ -147,6 +148,11 @@
bool LogMemory() { return log_memory_; }
Rendezvous* GetRendezvous() { return rendezvous_; }
+ std::unique_ptr<CollectiveExecutor::Handle> GetCollectiveExecutorHandle() {
+ return std::unique_ptr<CollectiveExecutor::Handle>(
+ new CollectiveExecutor::Handle(
+ collective_executor_mgr_->FindOrCreate(0), true /*inherit_ref*/));
+ }
const tensorflow::DeviceMgr* local_device_mgr() const {
return (local_device_manager_ != nullptr) ? local_device_manager_.get()
@@ -273,6 +279,8 @@
Env* const env_;
+ std::unique_ptr<CollectiveExecutorMgrInterface> collective_executor_mgr_;
+
#ifndef __ANDROID__
void CloseRemoteContexts();
diff --git a/tensorflow/core/common_runtime/eager/execute.cc b/tensorflow/core/common_runtime/eager/execute.cc
index 5bf7888..783baa9 100644
--- a/tensorflow/core/common_runtime/eager/execute.cc
+++ b/tensorflow/core/common_runtime/eager/execute.cc
@@ -263,7 +263,8 @@
// Note that it is not ideal, but currently ok, to set this
// attribute after computing the kernel cache key above.
if (op->is_function() && device != nullptr &&
- device->device_type() == "TPU") {
+ (device->device_type() == "TPU" || device->device_type() == "XLA_GPU" ||
+ device->device_type() == "XLA_CPU")) {
op->MutableAttrs()->Set(kXlaCompileAttr, true);
}
@@ -284,7 +285,8 @@
"Unable to find a FunctionLibraryRuntime corresponding to device ",
device->name());
}
- kernel = new KernelAndDevice(ctx->GetRendezvous(), ctx->LogMemory());
+ kernel = new KernelAndDevice(ctx->GetRendezvous(), ctx->LogMemory(),
+ ctx->GetCollectiveExecutorHandle());
status = KernelAndDevice::Init(ndef, flr, ctx->runner(), kernel);
if (!status.ok()) {
delete kernel;
diff --git a/tensorflow/core/common_runtime/eager/kernel_and_device.cc b/tensorflow/core/common_runtime/eager/kernel_and_device.cc
index 192d22d..317e9a1 100644
--- a/tensorflow/core/common_runtime/eager/kernel_and_device.cc
+++ b/tensorflow/core/common_runtime/eager/kernel_and_device.cc
@@ -84,6 +84,15 @@
tensorflow::HOST_MEMORY);
}
+ gtl::InlinedVector<DeviceContext*, 4> input_device_contexts;
+ for (int i = 0; i < inputs->size(); i++) {
+ DeviceContext* device_context = nullptr;
+ if (device_->tensorflow_gpu_device_info() != nullptr) {
+ device_context = device_->tensorflow_gpu_device_info()->default_context;
+ }
+ input_device_contexts.push_back(device_context);
+ }
+
OpKernelContext::Params params;
params.device = device_;
params.frame_iter = FrameAndIter(0, 0);
@@ -110,6 +119,9 @@
}
params.step_container = step_container;
+ params.collective_executor =
+ collective_executor_ ? collective_executor_->get() : nullptr;
+ params.input_device_contexts = &input_device_contexts;
OpKernelContext context(¶ms);
diff --git a/tensorflow/core/common_runtime/eager/kernel_and_device.h b/tensorflow/core/common_runtime/eager/kernel_and_device.h
index 52dac94..ee430b7 100644
--- a/tensorflow/core/common_runtime/eager/kernel_and_device.h
+++ b/tensorflow/core/common_runtime/eager/kernel_and_device.h
@@ -23,6 +23,7 @@
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
@@ -55,10 +56,16 @@
KernelAndDevice* out);
KernelAndDevice(tensorflow::Rendezvous* rendez, bool log_memory)
+ : KernelAndDevice(rendez, log_memory, nullptr) {}
+
+ KernelAndDevice(
+ tensorflow::Rendezvous* rendez, bool log_memory,
+ std::unique_ptr<CollectiveExecutor::Handle> collective_executor)
: device_(nullptr),
flr_(nullptr),
rendez_(rendez),
- log_memory_(log_memory) {}
+ log_memory_(log_memory),
+ collective_executor_(std::move(collective_executor)) {}
// TODO(ashankar): Handle list-valued inputs.
Status Run(std::vector<Tensor>* inputs, std::vector<Tensor>* outputs,
@@ -92,6 +99,7 @@
std::function<void(std::function<void()>)>* runner_;
std::function<void(std::function<void()>)> default_runner_;
const bool log_memory_;
+ const std::unique_ptr<CollectiveExecutor::Handle> collective_executor_;
};
} // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/tensor_handle.cc b/tensorflow/core/common_runtime/eager/tensor_handle.cc
index d8d6b7a..0acd160 100644
--- a/tensorflow/core/common_runtime/eager/tensor_handle.cc
+++ b/tensorflow/core/common_runtime/eager/tensor_handle.cc
@@ -184,10 +184,7 @@
bool is_same_device = (srcd == dstd) || (srcd->name() == dstd->name());
const bool dst_cpu = dstd->tensorflow_gpu_device_info() == nullptr;
const bool src_cpu = srcd->tensorflow_gpu_device_info() == nullptr;
- // both_on_cpu can be true and yet is_same_device is false, if one of src/dst
- // has device type XLA_CPU, and the other CPU.
- const bool both_on_cpu = src_cpu && dst_cpu;
- if (is_same_device || both_on_cpu) {
+ if (is_same_device) {
*output = new tensorflow::TensorHandle(*src, dstd, dstd, ctx);
return tensorflow::Status::OK();
}
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc
index 5152d97..14b57cc 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc
@@ -601,7 +601,9 @@
[to, copy](StatusCallback done_,
// Begin unbound arguments.
const Status& s) {
- *to = std::move(*copy);
+ if (s.ok()) {
+ *to = std::move(*copy);
+ }
delete copy;
done_(s);
},
diff --git a/tensorflow/core/common_runtime/gpu/gpu_process_state.cc b/tensorflow/core/common_runtime/gpu/gpu_process_state.cc
index a9a19f0..8167cfb 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_process_state.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_process_state.cc
@@ -70,7 +70,10 @@
// Return the NUMA node associated with the GPU's StreamExecutor.
se::StreamExecutor* se =
GpuIdUtil::ExecutorForTfGpuId(tf_gpu_id).ValueOrDie();
- return se->GetDeviceDescription().numa_node();
+ int numa_node = se->GetDeviceDescription().numa_node();
+ // bus_id must be non-negative. If the numa_node is not known,
+ // use 0.
+ return numa_node >= 0 ? numa_node : 0;
}
Allocator* GPUProcessState::GetGPUAllocator(const GPUOptions& options,
@@ -97,6 +100,7 @@
PlatformGpuId platform_gpu_id;
TF_CHECK_OK(GpuIdManager::TfToPlatformGpuId(tf_gpu_id, &platform_gpu_id));
int bus_id = BusIdForGPU(tf_gpu_id);
+ DCHECK_GE(bus_id, 0);
while (bus_id >= gpu_visitors_.size()) {
gpu_visitors_.push_back({});
}
@@ -249,6 +253,7 @@
CHECK(gpu_allocators_.empty()) // Crash OK
<< "AddGPUAllocVisitor must be called before "
"first call to GetGPUAllocator.";
+ DCHECK_GE(bus_id, 0);
while (bus_id >= static_cast<int64>(gpu_visitors_.size())) {
gpu_visitors_.push_back(std::vector<SubAllocator::Visitor>());
}
diff --git a/tensorflow/core/common_runtime/gpu/gpu_util_platform_specific.cc b/tensorflow/core/common_runtime/gpu/gpu_util_platform_specific.cc
index 4bc88ff..0ef39fb 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_util_platform_specific.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_util_platform_specific.cc
@@ -37,6 +37,14 @@
GPUUtil::CopyGPUTensorToCPU(device, this, device_tensor, cpu_tensor, done);
}
+void GPUDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,
+ Device* device,
+ Tensor* output_tensor,
+ StatusCallback done) const {
+ GPUUtil::CopyGPUTensorToSameGPU(device, this, input_tensor, output_tensor,
+ done);
+}
+
Status GPUDeviceContext::ThenExecute(Device* device, se::Stream* stream,
std::function<void()> func) {
const DeviceBase::GpuDeviceInfo* gpu_info =
diff --git a/tensorflow/core/common_runtime/gpu_device_context.h b/tensorflow/core/common_runtime/gpu_device_context.h
index 3603808..f513526 100644
--- a/tensorflow/core/common_runtime/gpu_device_context.h
+++ b/tensorflow/core/common_runtime/gpu_device_context.h
@@ -57,6 +57,10 @@
Device* device, Tensor* cpu_tensor,
StatusCallback done) override;
+ void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
+ Tensor* output_tensor,
+ StatusCallback done) const override;
+
void MaintainLifetimeOnStream(const Tensor* t,
se::Stream* stream) const override {}
diff --git a/tensorflow/core/common_runtime/graph_execution_state.cc b/tensorflow/core/common_runtime/graph_execution_state.cc
index 880806f..04d658f 100644
--- a/tensorflow/core/common_runtime/graph_execution_state.cc
+++ b/tensorflow/core/common_runtime/graph_execution_state.cc
@@ -546,10 +546,6 @@
std::unique_ptr<Graph> new_graph(new Graph(OpRegistry::Global()));
GraphConstructorOptions opts;
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(opts, *graph_def, new_graph.get()));
- for (const Node* n : new_graph->nodes()) {
- VLOG(2) << "Mapping " << n->name() << " to " << n->cost_id();
- node_name_to_cost_id_map_[n->name()] = n->cost_id();
- }
if (session_options_ &&
session_options_->config.graph_options().place_pruned_graph()) {
// Rewrite the graph before placement.
@@ -578,6 +574,11 @@
TF_RETURN_IF_ERROR(OptimizationPassRegistry::Global()->RunGrouping(
OptimizationPassRegistry::POST_PLACEMENT, optimization_options));
+ for (const Node* n : new_graph->nodes()) {
+ VLOG(2) << "Mapping " << n->name() << " to " << n->cost_id();
+ node_name_to_cost_id_map_[n->name()] = n->cost_id();
+ }
+
SaveStatefulNodes(new_graph.get());
graph_ = new_graph.release();
return Status::OK();
diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc
index bc8ba6e..59bb18e 100644
--- a/tensorflow/core/distributed_runtime/master_session.cc
+++ b/tensorflow/core/distributed_runtime/master_session.cc
@@ -1352,7 +1352,9 @@
&workers[i].call_opts, &workers[i].request, &workers[i].response, cb);
}
- done.Wait();
+ if (!done.WaitFor(std::chrono::milliseconds(10000))) {
+ LOG(WARNING) << "Timeout for closing worker session";
+ }
for (size_t i = 0; i < workers.size(); ++i) {
status.Update(workers[i].status);
}
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_channel.cc b/tensorflow/core/distributed_runtime/rpc/grpc_channel.cc
index 781b7d6..e5634d3 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_channel.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_channel.cc
@@ -60,6 +60,8 @@
// TODO(mrry): Implement secure channels.
::grpc::ChannelArguments args;
args.SetInt(GRPC_ARG_MAX_MESSAGE_LENGTH, std::numeric_limits<int32>::max());
+ args.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, std::numeric_limits<int>::max());
+ args.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, std::numeric_limits<int>::max());
// NOTE(mrry): Some versions of gRPC use a 20-second minimum backoff
// on connection failure, which makes our tests time out.
args.SetInt("grpc.testing.fixed_reconnect_backoff_ms", 1000);
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
index cbd5cd9..0851860 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
@@ -110,6 +110,8 @@
// - worker_env_.compute_pool
}
+void GrpcServer::MaybeMutateBuilder(::grpc::ServerBuilder* builder) {}
+
Status GrpcServer::Init(
ServiceInitFunction service_func,
const RendezvousMgrCreationFunction& rendezvous_mgr_func,
@@ -191,6 +193,11 @@
builder.AddListeningPort(strings::StrCat("0.0.0.0:", requested_port),
GetServerCredentials(server_def_), &bound_port_);
builder.SetMaxMessageSize(std::numeric_limits<int32>::max());
+ builder.AddChannelArgument(GRPC_ARG_KEEPALIVE_TIME_MS,
+ std::numeric_limits<int>::max());
+ builder.AddChannelArgument(GRPC_ARG_KEEPALIVE_TIMEOUT_MS,
+ std::numeric_limits<int>::max());
+
builder.SetOption(
std::unique_ptr<::grpc::ServerBuilderOption>(new NoReusePortOption));
// Allow subclasses to specify more args to pass to the gRPC server.
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h
index c1395ab..c7f543e 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h
@@ -62,7 +62,7 @@
GrpcServer(const ServerDef& server_def, Env* env);
// Allow children classes to override this and provide custom args to the
// server before it is constructed. Default behavior is to do nothing.
- virtual void MaybeMutateBuilder(::grpc::ServerBuilder* builder) {}
+ virtual void MaybeMutateBuilder(::grpc::ServerBuilder* builder);
public:
static Status Create(const ServerDef& server_def, Env* env,
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_session.h b/tensorflow/core/distributed_runtime/rpc/grpc_session.h
index e00cf97..a3ed3ec 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_session.h
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_session.h
@@ -111,7 +111,7 @@
LOCKS_EXCLUDED(mu_);
private:
- SessionOptions options_;
+ const SessionOptions options_;
std::unique_ptr<MasterInterface> master_;
mutex mu_;
diff --git a/tensorflow/core/framework/dataset.cc b/tensorflow/core/framework/dataset.cc
index 6af1415..6e21433 100644
--- a/tensorflow/core/framework/dataset.cc
+++ b/tensorflow/core/framework/dataset.cc
@@ -17,6 +17,8 @@
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/variant_encode_decode.h"
+#include "tensorflow/core/framework/variant_op_registry.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/platform/mutex.h"
@@ -74,6 +76,113 @@
DatasetBase* const dataset_; // Owns one reference.
};
+const char kWrappedDatasetVariantTypeName[] =
+ "tensorflow::data::WrappedDatasetVariant";
+
+class WrappedDatasetVariantWrapper {
+ public:
+ WrappedDatasetVariantWrapper() {}
+
+ explicit WrappedDatasetVariantWrapper(const Tensor& ds_tensor)
+ : ds_tensor_(ds_tensor) {}
+
+ Tensor get() const { return ds_tensor_; }
+
+ string TypeName() const { return "tensorflow::WrappedDatasetVariantWrapper"; }
+
+ string DebugString() const {
+ return "tensorflow::WrappedDatasetVariantWrapper::DebugString";
+ }
+
+ void Encode(VariantTensorData* data) const {
+ *(data->add_tensors()) = ds_tensor_;
+ }
+
+ bool Decode(const VariantTensorData& data) {
+ ds_tensor_ = data.tensors(0);
+ return true;
+ }
+
+ private:
+ Tensor ds_tensor_;
+};
+
+class WrapDatasetVariantOp : public OpKernel {
+ public:
+ explicit WrapDatasetVariantOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+
+ void Compute(OpKernelContext* ctx) override {
+ const Tensor& tensor = ctx->input(0);
+ OP_REQUIRES(ctx,
+ tensor.dtype() == DT_VARIANT &&
+ TensorShapeUtils::IsScalar(tensor.shape()),
+ errors::InvalidArgument(
+ "Dataset tensor must be a scalar of dtype DT_VARIANT."));
+ DatasetBase* unused;
+ OP_REQUIRES_OK(ctx, GetDatasetFromVariantTensor(tensor, &unused));
+ Tensor* output = nullptr;
+ OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &output));
+ output->scalar<Variant>()() = WrappedDatasetVariantWrapper(tensor);
+ }
+};
+
+REGISTER_KERNEL_BUILDER(Name("WrapDatasetVariant").Device(DEVICE_CPU),
+ WrapDatasetVariantOp);
+REGISTER_KERNEL_BUILDER(Name("WrapDatasetVariant")
+ .HostMemory("input_handle")
+ .HostMemory("output_handle")
+ .Device(DEVICE_GPU),
+ WrapDatasetVariantOp);
+
+class UnwrapDatasetVariantOp : public OpKernel {
+ public:
+ explicit UnwrapDatasetVariantOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+
+ void Compute(OpKernelContext* ctx) override {
+ const Tensor& tensor = ctx->input(0);
+ OP_REQUIRES(ctx,
+ tensor.dtype() == DT_VARIANT &&
+ TensorShapeUtils::IsScalar(tensor.shape()),
+ errors::InvalidArgument(
+ "Dataset tensor must be a scalar of dtype DT_VARIANT."));
+ Variant variant = tensor.scalar<Variant>()();
+ const WrappedDatasetVariantWrapper* wrapper =
+ variant.get<WrappedDatasetVariantWrapper>();
+ OP_REQUIRES(ctx, wrapper != nullptr,
+ errors::InvalidArgument(
+ "Tensor must be a WrappedDataset variant object."));
+ Tensor ds_tensor = wrapper->get();
+ OP_REQUIRES_OK(ctx, ctx->set_output("output_handle", ds_tensor));
+ }
+};
+
+REGISTER_KERNEL_BUILDER(Name("UnwrapDatasetVariant").Device(DEVICE_CPU),
+ UnwrapDatasetVariantOp);
+REGISTER_KERNEL_BUILDER(Name("UnwrapDatasetVariant")
+ .HostMemory("input_handle")
+ .HostMemory("output_handle")
+ .Device(DEVICE_GPU),
+ UnwrapDatasetVariantOp);
+
+static Status WrappedDatasetVariantDeviceCopy(
+ const WrappedDatasetVariantWrapper& from, WrappedDatasetVariantWrapper* to,
+ const UnaryVariantOpRegistry::AsyncTensorDeviceCopyFn& copy) {
+ *to = WrappedDatasetVariantWrapper(from);
+ return Status::OK();
+}
+
+#define REGISTER_OPTIONAL_COPY(DIRECTION) \
+ INTERNAL_REGISTER_UNARY_VARIANT_DEVICE_COPY_FUNCTION( \
+ WrappedDatasetVariantWrapper, DIRECTION, \
+ WrappedDatasetVariantDeviceCopy)
+
+REGISTER_OPTIONAL_COPY(VariantDeviceCopyDirection::HOST_TO_DEVICE);
+REGISTER_OPTIONAL_COPY(VariantDeviceCopyDirection::DEVICE_TO_HOST);
+REGISTER_OPTIONAL_COPY(VariantDeviceCopyDirection::DEVICE_TO_DEVICE);
+
+REGISTER_UNARY_VARIANT_DECODE_FUNCTION(WrappedDatasetVariantWrapper,
+ kWrappedDatasetVariantTypeName);
+
} // namespace
Status GraphDefBuilderWrapper::AddDataset(
@@ -206,6 +315,20 @@
return HasAttr(op_def, attr_name);
}
+int64 GetAllocatedBytes(const std::vector<Tensor>& element) {
+ int64 allocated_bytes = 0;
+ DatasetBase* dataset;
+ for (auto& tensor : element) {
+ if (tensor.dtype() == DT_VARIANT &&
+ GetDatasetFromVariantTensor(tensor, &dataset).ok()) {
+ allocated_bytes += dataset->AllocatedBytes();
+ } else {
+ allocated_bytes += tensor.AllocatedBytes();
+ }
+ }
+ return allocated_bytes;
+}
+
Status GetDatasetFromVariantTensor(const Tensor& tensor,
DatasetBase** out_dataset) {
if (!(tensor.dtype() == DT_VARIANT &&
diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index a00768f..cca10fa 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -50,8 +50,6 @@
class Node;
namespace data {
-// A constant that can be used to enable auto-tuning.
-constexpr int kAutoTune = -1;
constexpr int kInfiniteCardinality = -1;
constexpr int kUnknownCardinality = -2;
@@ -285,6 +283,7 @@
function_library(ctx->function_library()),
lib(ctx->lib()),
function_handle_cache(ctx->function_handle_cache()),
+ resource_mgr(ctx->resource_mgr()),
model(ctx->model()),
runner(*(ctx->runner())),
runner_threadpool_size(ctx->runner_threadpool_size()),
@@ -324,6 +323,10 @@
// A FunctionHandleCache that owns all the function handles. Not owned.
FunctionHandleCache* function_handle_cache = nullptr;
+ // A resource manager for storing dataset-related state, e.g. random
+ // seeds or cached tensors. Not owned.
+ ResourceMgr* resource_mgr = nullptr;
+
// If non-null, identifies the object used for performance modeling.
std::shared_ptr<model::Model> model = nullptr;
@@ -363,6 +366,8 @@
return params_.function_handle_cache;
}
+ ResourceMgr* resource_mgr() { return params_.resource_mgr; }
+
const std::shared_ptr<model::Model>& model() { return params_.model; }
std::function<void(std::function<void()>)>* runner() {
@@ -534,6 +539,25 @@
Params params_;
};
+// Returns the number of bytes allocated for the given tensor.
+int64 GetAllocatedBytes(const std::vector<Tensor>& element);
+
+// Validates and extracts a `DatasetBase` object from `tensor`.
+//
+// `tensor` must have been written by a call to SetVariantTensorToDataset().
+//
+// The retrieved pointer is a borrowed reference to the dataset, which is owned
+// by the tensor. The consumer must either acquire its own reference to the
+// dataset by calling `(*out_dataset)->Ref()`, or ensure that `tensor` is not
+// destroyed or mutated while the retrieved pointer is in use.
+Status GetDatasetFromVariantTensor(const Tensor& tensor,
+ DatasetBase** out_dataset);
+
+// Stores a `DatasetBase` object in `tensor`.
+//
+// The ownership of `dataset` is transferred to `tensor`.
+Status StoreDatasetInVariantTensor(DatasetBase* dataset, Tensor* tensor);
+
// Represents a (potentially infinite) range of outputs, where each
// output is a tuple of tensors.
class DatasetBase : public core::RefCounted {
@@ -587,12 +611,15 @@
// in the outputs of this dataset.
virtual const std::vector<PartialTensorShape>& output_shapes() const = 0;
- // A human-readable debug string for this dataset.
- virtual string DebugString() const = 0;
+ // Returns the number of bytes allocated for tensors of this dataset.
+ virtual int64 AllocatedBytes() const { return 0; }
// Returns the cardinality of this dataset.
virtual int64 Cardinality() const { return kUnknownCardinality; }
+ // A human-readable debug string for this dataset.
+ virtual string DebugString() const = 0;
+
// Serializes the dataset and writes it to the `writer`.
virtual Status Save(SerializationContext* ctx,
IteratorStateWriter* writer) const;
@@ -694,18 +721,36 @@
return model::MakeUnknownNode(std::move(args));
}
- // When performance modeling is enabled, this method records the fact that
- // this iterator has produced an element.
+ // When modeling is enabled, this method records the fact that this iterator
+ // has dequeued an element from an internal buffer.
+ void RecordBufferDequeue(IteratorContext* ctx,
+ const std::vector<Tensor>& element) {
+ if (collect_resource_usage(ctx)) {
+ node_->add_buffered_bytes(-GetAllocatedBytes(element));
+ }
+ }
+
+ // When modeling is enabled, this method records the fact that this iterator
+ // has enqueued an element in an internal buffer.
+ void RecordBufferEnqueue(IteratorContext* ctx,
+ const std::vector<Tensor>& element) {
+ if (collect_resource_usage(ctx)) {
+ node_->add_buffered_bytes(GetAllocatedBytes(element));
+ }
+ }
+
+ // When modeling is enabled, this method records the fact that this iterator
+ // has produced an element.
void RecordElement(IteratorContext* ctx) {
if (node_) {
node_->record_element();
}
}
- // When performance modeling is enabled, this method records the fact that
- // a thread of this iterator has started work.
+ // When modeling is enabled, this method records the fact that a thread of
+ // this iterator has started work.
void RecordStart(IteratorContext* ctx, bool stop_output = false) {
- if (node_) {
+ if (collect_resource_usage(ctx)) {
int64 now_nanos = Env::Default()->NowNanos();
if (stop_output && node_->output()) {
node_->output()->record_stop(now_nanos);
@@ -714,10 +759,10 @@
}
}
- // When performance modeling is enabled, this method records the fact that
- // a thread of this iterator has stopped work.
+ // When modeling is enabled, this method records the fact that a thread of
+ // this iterator has stopped work.
void RecordStop(IteratorContext* ctx, bool start_output = false) {
- if (node_) {
+ if (collect_resource_usage(ctx)) {
int64 now_nanos = Env::Default()->NowNanos();
node_->record_stop(now_nanos);
if (start_output && node_->output()) {
@@ -727,6 +772,11 @@
}
private:
+ inline bool collect_resource_usage(IteratorContext* ctx) {
+ auto model = ctx->model();
+ return model && model->collect_resource_usage() && node_;
+ }
+
BaseParams params_;
};
@@ -826,22 +876,6 @@
DatasetBase** output) = 0;
};
-// Validates and extracts a `DatasetBase` object from `tensor`.
-//
-// `tensor` must have been written by a call to SetVariantTensorToDataset().
-//
-// The retrieved pointer is a borrowed reference to the dataset, which is owned
-// by the tensor. The consumer must either acquire its own reference to the
-// dataset by calling `(*out_dataset)->Ref()`, or ensure that `tensor` is not
-// destroyed or mutated while the retrieved pointer is in use.
-Status GetDatasetFromVariantTensor(const Tensor& tensor,
- DatasetBase** out_dataset);
-
-// Stores a `DatasetBase` object in `tensor`.
-//
-// The ownership of `dataset` is transferred to `tensor`.
-Status StoreDatasetInVariantTensor(DatasetBase* dataset, Tensor* tensor);
-
// A simple background worker that executes closures asynchronously and without
// blocking.
//
diff --git a/tensorflow/core/framework/device_base.h b/tensorflow/core/framework/device_base.h
index 446c31b..321947a 100644
--- a/tensorflow/core/framework/device_base.h
+++ b/tensorflow/core/framework/device_base.h
@@ -82,6 +82,13 @@
done(errors::Internal("Unrecognized device type in CPU-to-device Copy"));
}
+ // Copies a tensor in this device.
+ virtual void CopyTensorInSameDevice(const Tensor* input_tensor,
+ Device* device, Tensor* output_tensor,
+ StatusCallback done) const {
+ done(errors::Unimplemented("Copy in same device not implemented."));
+ }
+
// "device_tensor" is a tensor on a non-CPU device. Copies
// device_tensor into "cpu_tensor". "cpu_tensor" must be allocated
// to be of the same size as "device_tensor".
diff --git a/tensorflow/core/framework/model.cc b/tensorflow/core/framework/model.cc
index 3bd5b72..b7c6d80 100644
--- a/tensorflow/core/framework/model.cc
+++ b/tensorflow/core/framework/model.cc
@@ -356,6 +356,8 @@
if (output) {
output->add_input(node);
}
+ collect_resource_usage_ =
+ collect_resource_usage_ || node->has_tunable_parameters();
lookup_table_.insert(std::make_pair(name, node));
return node;
}
@@ -441,7 +443,7 @@
void Model::RecordStart(const string& name, bool stop_output) {
tf_shared_lock l(mu_);
auto node = gtl::FindOrNull(lookup_table_, name);
- if (node) {
+ if (collect_resource_usage_ && node) {
int64 now_nanos = Env::Default()->NowNanos();
if (stop_output && (*node)->output()) {
(*node)->output()->record_stop(now_nanos);
@@ -453,7 +455,7 @@
void Model::RecordStop(const string& name, bool start_output) {
tf_shared_lock l(mu_);
auto node = gtl::FindOrNull(lookup_table_, name);
- if (node) {
+ if (collect_resource_usage_ && node) {
int64 now_nanos = Env::Default()->NowNanos();
(*node)->record_stop(now_nanos);
if (start_output && (*node)->output()) {
diff --git a/tensorflow/core/framework/model.h b/tensorflow/core/framework/model.h
index 24aa563..c3a6942 100644
--- a/tensorflow/core/framework/model.h
+++ b/tensorflow/core/framework/model.h
@@ -34,18 +34,24 @@
namespace data {
namespace model {
+// A constant that can be used to enable auto-tuning.
+constexpr int kAutoTune = -1;
+
// Represents thread-safe state that can be shared between an input pipeline and
// the performance model.
struct SharedState {
public:
SharedState(int64 value, std::shared_ptr<mutex> mu,
std::shared_ptr<condition_variable> cond_var)
- : value(value), mu(std::move(mu)), cond_var(std::move(cond_var)) {}
+ : value(value),
+ mu(std::move(mu)),
+ cond_var(std::move(cond_var)),
+ tunable(value == kAutoTune) {}
int64 value;
std::shared_ptr<mutex> mu;
std::shared_ptr<condition_variable> cond_var;
- bool tunable = false;
+ const bool tunable;
};
// Represents a parameter.
@@ -112,6 +118,12 @@
explicit Node(Args args)
: id_(args.id), name_(args.name), output_(args.output.get()) {}
+ // Increments the bytes buffered by the given delta.
+ void add_buffered_bytes(int64 delta) LOCKS_EXCLUDED(mu_) {
+ mutex_lock l(mu_);
+ buffered_bytes_ += delta;
+ }
+
// Adds an input.
void add_input(std::shared_ptr<Node> node) LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
@@ -124,18 +136,33 @@
processing_time_ += delta;
}
+ // Returns the number of bytes stored in this node's buffer.
+ int64 buffered_bytes() const LOCKS_EXCLUDED(mu_) {
+ tf_shared_lock l(mu_);
+ return buffered_bytes_;
+ }
+
+ // Indicates whether the node has tunable parameters.
+ bool has_tunable_parameters() const LOCKS_EXCLUDED(mu_) {
+ tf_shared_lock l(mu_);
+ for (const auto& pair : parameters_) {
+ if (pair.second->state->tunable) return true;
+ }
+ return false;
+ }
+
// Returns the unique node ID.
int64 id() const LOCKS_EXCLUDED(mu_) { return id_; }
- // Returns the node name.
- const string& name() const { return name_; }
-
// Returns the node inputs.
std::list<std::shared_ptr<Node>> inputs() const LOCKS_EXCLUDED(mu_) {
tf_shared_lock l(mu_);
return inputs_;
}
+ // Returns the node name.
+ const string& name() const { return name_; }
+
// Returns the number of elements produced by the node.
int64 num_elements() const LOCKS_EXCLUDED(mu_) {
tf_shared_lock l(mu_);
@@ -185,7 +212,8 @@
// Collects tunable parameters in the subtree rooted in this node.
void CollectTunableParameters(
- std::vector<std::shared_ptr<Parameter>>* parameters) LOCKS_EXCLUDED(mu_) {
+ std::vector<std::shared_ptr<Parameter>>* parameters) const
+ LOCKS_EXCLUDED(mu_) {
tf_shared_lock l(mu_);
for (auto& pair : parameters_) {
if (pair.second->state->tunable) {
@@ -219,6 +247,7 @@
LOCKS_EXCLUDED(mu_) {
tf_shared_lock l(mu_);
std::shared_ptr<Node> result = Clone(output);
+ result->buffered_bytes_ = buffered_bytes_;
result->processing_time_ = processing_time_;
result->num_elements_ = num_elements_;
result->parameters_ = parameters_;
@@ -274,6 +303,7 @@
mutable mutex mu_;
const int64 id_;
const string name_;
+ int64 buffered_bytes_ GUARDED_BY(mu_) = 0;
int64 processing_time_ GUARDED_BY(mu_) = 0;
int64 num_elements_ GUARDED_BY(mu_) = 0;
std::map<std::thread::id, int64> work_start_ GUARDED_BY(mu_);
@@ -329,7 +359,10 @@
// implementation of `DatasetBase` and `DatasetBaseIterator` respectively.
class Model {
public:
- Model() = default;
+ Model() : collect_resource_usage_(false) {}
+
+ // Indicates whether to collect resource usage.
+ bool collect_resource_usage() const { return collect_resource_usage_; }
// Adds a node with the given name and given output.
std::shared_ptr<Node> AddNode(Node::Factory factory, const string& name,
@@ -373,6 +406,14 @@
int64 id_counter_ GUARDED_BY(mu_) = 1;
std::shared_ptr<Node> output_ GUARDED_BY(mu_);
std::map<string, std::shared_ptr<Node>> lookup_table_ GUARDED_BY(mu_);
+
+ // Indicates whether the modeling framework should collect resource usage
+ // (e.g. CPU, memory). The logic for collecting this information assumes that
+ // the collection is not repeatedly disabled and enabled. As a consequence,
+ // the implementation starts collecting resource usage when it encounters a
+ // tunable parameter (because the information is used for for tuning the value
+ // of the parameter) and never stops.
+ std::atomic<bool> collect_resource_usage_;
};
} // namespace model
diff --git a/tensorflow/core/framework/model_test.cc b/tensorflow/core/framework/model_test.cc
index 53e35f2..90bd570 100644
--- a/tensorflow/core/framework/model_test.cc
+++ b/tensorflow/core/framework/model_test.cc
@@ -330,6 +330,62 @@
EXPECT_EQ(100, unknown->OutputTime(&input_times));
}
+class TestNode : public model::Node {
+ public:
+ using model::Node::Node;
+
+ virtual ~TestNode() {}
+
+ protected:
+ std::shared_ptr<Node> Clone(std::shared_ptr<Node> output) const override
+ SHARED_LOCKS_REQUIRED(mu_) {
+ return nullptr;
+ }
+
+ int64 OutputTimeLocked(std::vector<int64>* input_times) const override
+ SHARED_LOCKS_REQUIRED(mu_) {
+ return 0;
+ }
+
+ int64 ProcessingTimeLocked() const override SHARED_LOCKS_REQUIRED(mu_) {
+ return 0;
+ }
+};
+
+TEST(SetterGetterTest, Node) {
+ std::shared_ptr<TestNode> node =
+ std::make_shared<TestNode>(model::Node::Args{-1, "TestNode", nullptr});
+ EXPECT_EQ(-1, node->id());
+ EXPECT_EQ("TestNode", node->name());
+ EXPECT_EQ(nullptr, node->output());
+
+ EXPECT_EQ(0, node->buffered_bytes());
+ node->add_buffered_bytes(42);
+ EXPECT_EQ(42, node->buffered_bytes());
+
+ EXPECT_EQ(0, node->processing_time());
+ node->record_start(1);
+ EXPECT_EQ(0, node->processing_time());
+ node->record_stop(41);
+ EXPECT_EQ(40, node->processing_time());
+ node->add_processing_time(2);
+ EXPECT_EQ(42, node->processing_time());
+
+ std::shared_ptr<TestNode> input =
+ std::make_shared<TestNode>(model::Node::Args{-1, "TestInput", node});
+ EXPECT_EQ(node.get(), input->output());
+ EXPECT_EQ(0, node->inputs().size());
+ node->add_input(input);
+ EXPECT_EQ(1, node->inputs().size());
+ EXPECT_EQ(input, node->inputs().front());
+ node->remove_input(input);
+ EXPECT_EQ(0, node->inputs().size());
+
+ EXPECT_EQ(0, node->num_elements());
+ node->record_element();
+ EXPECT_EQ(1, node->num_elements());
+}
+
} // namespace
} // namespace model
} // namespace data
diff --git a/tensorflow/core/framework/op_kernel.h b/tensorflow/core/framework/op_kernel.h
index 9f4c57e..19a0c5e 100644
--- a/tensorflow/core/framework/op_kernel.h
+++ b/tensorflow/core/framework/op_kernel.h
@@ -1527,6 +1527,7 @@
template <typename T>
T* OpKernelContext::input_device_context(int index) {
+ DCHECK_NE(params_->input_device_contexts, nullptr);
DCHECK_GE(index, 0);
DCHECK_LT(index, params_->input_device_contexts->size());
static_assert(std::is_base_of<DeviceContext, T>::value,
@@ -1535,6 +1536,7 @@
}
inline DeviceContext* OpKernelContext::input_device_context(int index) {
+ DCHECK_NE(params_->input_device_contexts, nullptr);
DCHECK_GE(index, 0);
DCHECK_LT(index, params_->input_device_contexts->size());
return (*params_->input_device_contexts)[index];
diff --git a/tensorflow/core/framework/rendezvous_test.cc b/tensorflow/core/framework/rendezvous_test.cc
index de148f0..7a777f0 100644
--- a/tensorflow/core/framework/rendezvous_test.cc
+++ b/tensorflow/core/framework/rendezvous_test.cc
@@ -278,6 +278,12 @@
~DummyDeviceContext() override {}
int stream_id() const { return stream_id_; }
+ void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
+ Tensor* output_tensor,
+ StatusCallback done) const override {
+ done(Status::OK());
+ }
+
private:
const int stream_id_;
};
diff --git a/tensorflow/core/framework/resource_var.h b/tensorflow/core/framework/resource_var.h
index ff7b3e7..f5de5db 100644
--- a/tensorflow/core/framework/resource_var.h
+++ b/tensorflow/core/framework/resource_var.h
@@ -20,14 +20,46 @@
namespace tensorflow {
-// Resource stored by variables in the resource manager
-// (new, resource-style version).
+// Resource stored by variables in the resource manager (new, resource-style
+// version).
+//
+// These variables have a mixed access mode: they can operate on copy-on-write
+// mode (the default) or copy-on-read mode (used only for sparse access).
+//
+// When copy-on-write mode is enabled reading the value of the variable involves
+// grabbing its mutex in shared mode and aliasing the internal tensor as the
+// output of the read operation, increasing its reference count. Writing,
+// conversely, works by, under an exclusive lock, detecting whether there are
+// outstanding aliases of the tensor, using the reference count, copying the
+// tensor if they exist, and writing to either the original or a copy with no
+// outstanding aliases. Sparse operations are not supported in copy-on-write
+// mode.
+//
+// When a variable is accessed sparsely it switches to copy-on-read mode. To
+// switch we need to grab an exclusive lock and might (if there are aliases)
+// need to copy the entire tensor. Once copy-on-read mode is enabled, no tensor
+// is allowed to alias the variable's internal tensor. This means dense reads
+// must return a copy of the variable, done while holding a shared lock. Dense
+// writes do not need to check whether aliases exist, and can always write
+// directly to the buffer without making a copy, while holding an exclusive
+// lock. Sparse reads and sparse writes, on the other hand, can be done under a
+// shared or exclusive mutex (the damage from writes under a shared mutex is
+// limited since no other buffer is allowed to alias the variable's
+// buffer). Using an exclusive mutex disallows concurrent writes and concurrent
+// sparse reads, providing some extra safety at the expense of performance,
+// while shared mutex allow for "hogwild" behavior. Doing sparse writes under a
+// shared mutex prevents them from overlapping with dense writes, which is
+// necessary as dense writes can change the shape the of the tensor.
+//
+// Transitioning a variable from copy-on-read mode to copy-on-write mode is
+// currently not supported. To upgrade a variable from copy-on-write to
+// copy-on-read use `EnsureSparseVariableAccess()`, and then grab the variable's
+// mutex as desired. To access the variable in dense mode grab the mutex either
+// directly or via `MaybeLockVariableInputMutexesInOrder` on all variables being
+// modified and then call `PrepareToUpdateVariable` on them in any order.
class Var : public ResourceBase {
public:
explicit Var(DataType dtype) : tensor_(dtype) {}
- // Not copyable or movable.
- Var(const Var&) = delete;
- Var& operator=(const Var&) = delete;
// When locking multiple variables, the locks must be acquired in order of
// increasing mu() address.
@@ -48,11 +80,19 @@
bool is_initialized = false; // GUARDED_BY(mu_) but annotalysis doesn't like
// it.
+ // Also fake-guarded by mu_. Should be set to True whenever any sparse
+ // operation uses the variable. Once this is true no tensor is allowed to
+ // alias the memory of the variable, and we always copy the variable on
+ // reads. This allows sparse operations to happen with only a shared lock if
+ // so desired.
+ std::atomic<bool> copy_on_read_mode{false};
+
private:
mutex mu_;
Tensor tensor_;
~Var() override {}
+ TF_DISALLOW_COPY_AND_ASSIGN(Var);
};
} // end namespace tensorflow
diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h
index 6e03cf9..009dd08 100644
--- a/tensorflow/core/framework/tensor.h
+++ b/tensorflow/core/framework/tensor.h
@@ -45,6 +45,7 @@
class TensorCApi;
class TensorDescription;
class TensorProto;
+class Var;
namespace batch_util {
Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index);
@@ -581,11 +582,16 @@
friend class XlaTensor; // For access to RefCountIsOne().
friend class XlaTensorBuffer; // For access to the private constructor taking
// the buffer
+ friend class Var;
template <typename Device, typename T>
friend class AssignVariableOp; // For access to RefCountIsOne().
template <typename Device, typename T>
friend Status PrepareToUpdateVariable(
- OpKernelContext* ctx, Tensor* tensor); // For access to RefCountIsOne().
+ OpKernelContext* ctx, Tensor* tensor,
+ bool copy_on_read_mode); // For access to RefCountIsOne().
+ template <typename Device, typename T>
+ friend Status EnsureSparseVariableAccess(
+ OpKernelContext* ctx, Var* var); // For access to RefCountIsOne().
friend Status batch_util::CopyElementToSlice(
Tensor element, Tensor* parent,
int64 index); // For access to RefCountIsOne().
diff --git a/tensorflow/core/graph/graph.cc b/tensorflow/core/graph/graph.cc
index 550e3ef..223fc85 100644
--- a/tensorflow/core/graph/graph.cc
+++ b/tensorflow/core/graph/graph.cc
@@ -555,7 +555,13 @@
dst->DebugString());
}
TF_RETURN_IF_ERROR(IsValidOutputTensor(new_src, new_src_index));
- int dst_index = dst->in_edges().size();
+ // Find the current number of data inputs. We'll add the new edge to the next
+ // missing data input.
+ int dst_index = 0;
+ for (const Edge* edge : dst->in_edges()) {
+ if (edge->IsControlEdge()) continue;
+ ++dst_index;
+ }
TF_RETURN_IF_ERROR(IsValidInputTensor(dst, dst_index));
AddEdge(new_src, new_src_index, dst, dst_index);
dst->MaybeCopyOnWrite();
diff --git a/tensorflow/core/graph/graph_test.cc b/tensorflow/core/graph/graph_test.cc
index 333c325..e7762fd 100644
--- a/tensorflow/core/graph/graph_test.cc
+++ b/tensorflow/core/graph/graph_test.cc
@@ -799,5 +799,44 @@
BENCHMARK(BM_GraphCreation)->ArgPair(1 << 12, 16);
BENCHMARK(BM_GraphCreation)->ArgPair(1 << 15, 16);
+static void BM_ToGraphDef(int iters, int num_nodes, int num_edges_per_node) {
+ testing::StopTiming();
+ const GraphDef graph_def = CreateGraphDef(num_nodes, num_edges_per_node);
+ const auto registry = OpRegistry::Global();
+ GraphConstructorOptions opts;
+ // Warmup step.
+ Graph graph(registry);
+ TF_CHECK_OK(ConvertGraphDefToGraph(opts, graph_def, &graph));
+ int64 sum = 0;
+ testing::StartTiming();
+ for (int i = 0; i < iters; ++i) {
+ GraphDef graph_def;
+ graph.ToGraphDef(&graph_def);
+ sum += graph_def.node_size();
+ }
+ VLOG(1) << sum;
+ testing::StopTiming();
+}
+BENCHMARK(BM_ToGraphDef)->ArgPair(10, 2);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 6, 2);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 9, 2);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 12, 2);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 15, 2);
+BENCHMARK(BM_ToGraphDef)->ArgPair(10, 4);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 6, 4);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 9, 4);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 12, 4);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 15, 4);
+BENCHMARK(BM_ToGraphDef)->ArgPair(10, 8);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 6, 8);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 9, 8);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 12, 8);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 15, 8);
+BENCHMARK(BM_ToGraphDef)->ArgPair(10, 16);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 6, 16);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 9, 16);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 12, 16);
+BENCHMARK(BM_ToGraphDef)->ArgPair(1 << 15, 16);
+
} // namespace
} // namespace tensorflow
diff --git a/tensorflow/core/graph/mkl_layout_pass.cc b/tensorflow/core/graph/mkl_layout_pass.cc
index 7b12f0a..3c868dc 100644
--- a/tensorflow/core/graph/mkl_layout_pass.cc
+++ b/tensorflow/core/graph/mkl_layout_pass.cc
@@ -260,7 +260,10 @@
csinfo_.conv3d_grad_filter = "Conv3DBackpropFilterV2";
csinfo_.fused_batch_norm = "FusedBatchNorm";
csinfo_.fused_batch_norm_grad = "FusedBatchNormGrad";
+ csinfo_.fused_conv2d = "_FusedConv2D";
csinfo_.identity = "Identity";
+ csinfo_.leakyrelu = "LeakyRelu";
+ csinfo_.leakyrelu_grad = "LeakyReluGrad";
csinfo_.lrn = "LRN";
csinfo_.lrn_grad = "LRNGrad";
csinfo_.matmul = "MatMul";
@@ -274,6 +277,10 @@
csinfo_.mkl_conv2d_with_bias = "_MklConv2DWithBias";
csinfo_.mkl_conv2d_grad_filter_with_bias =
"_MklConv2DBackpropFilterWithBias";
+ csinfo_.mkl_fused_conv2d = "_MklFusedConv2D";
+ csinfo_.mkl_pad_with_conv2d = "_MklPadWithConv2D";
+ csinfo_.pad = "Pad";
+ csinfo_.pad_with_conv2d = "__MklDummyPadWithConv2D";
// Temporarily don't convert quantized operators into MKL versions for now.
// TODO(Intel-tf) Once all the relevant PRs have been merged then remove
// the ifdef.
@@ -377,6 +384,8 @@
{csinfo_.fused_batch_norm_grad,
mkl_op_registry::GetMklOpName(csinfo_.fused_batch_norm_grad),
CopyAttrsFusedBatchNorm, AlwaysRewrite});
+ rinfo_.push_back({csinfo_.fused_conv2d, csinfo_.mkl_fused_conv2d,
+ CopyAttrsFusedConv2D, FusedConv2DRewrite});
rinfo_.push_back({csinfo_.identity,
mkl_op_registry::GetMklOpName(csinfo_.identity),
CopyAttrsDataType, AlwaysRewrite});
@@ -385,6 +394,12 @@
rinfo_.push_back({csinfo_.lrn_grad,
mkl_op_registry::GetMklOpName(csinfo_.lrn_grad),
CopyAttrsLRN, LrnGradRewrite});
+ rinfo_.push_back({csinfo_.leakyrelu,
+ mkl_op_registry::GetMklOpName(csinfo_.leakyrelu),
+ CopyAttrsLeakyRelu, LeakyReluRewrite});
+ rinfo_.push_back({csinfo_.leakyrelu_grad,
+ mkl_op_registry::GetMklOpName(csinfo_.leakyrelu_grad),
+ CopyAttrsLeakyRelu, LeakyReluRewrite});
rinfo_.push_back({csinfo_.max_pool,
mkl_op_registry::GetMklOpName(csinfo_.max_pool),
CopyAttrsPooling, NonDepthBatchWisePoolRewrite});
@@ -402,6 +417,8 @@
CopyAttrsDataType, AlwaysRewrite});
rinfo_.push_back({csinfo_.mul, mkl_op_registry::GetMklOpName(csinfo_.mul),
CopyAttrsDataType, AlwaysRewrite});
+ rinfo_.push_back({csinfo_.pad_with_conv2d, csinfo_.mkl_pad_with_conv2d,
+ CopyAttrsPadWithConv2D, AlwaysRewrite});
#ifdef INTEL_MKL_QUANTIZED
rinfo_.push_back({csinfo_.quantized_avg_pool,
mkl_op_registry::GetMklOpName(csinfo_.quantized_avg_pool),
@@ -512,6 +529,10 @@
minfo_.push_back({csinfo_.conv2d_grad_filter, csinfo_.bias_add_grad,
csinfo_.conv2d_grad_filter_with_bias,
GetConv2DBackpropFilterOrBiasAddGrad});
+ minfo_.push_back(
+ {csinfo_.pad, csinfo_.conv2d, csinfo_.pad_with_conv2d, GetPadOrConv2D});
+ // Merge Pad and Conv2d, only if the pad op is "Pad"
+ // Doesn't merge if pad op is "PadV2" or "MirrorPad"
// The fusion patterns in "finfo_" that show up first will get applied
// first, for example, graph "A->B->C-D" and finfo_ is {A->B->C to ABC,
@@ -656,7 +677,10 @@
string conv3d_grad_filter;
string fused_batch_norm;
string fused_batch_norm_grad;
+ string fused_conv2d;
string identity;
+ string leakyrelu;
+ string leakyrelu_grad;
string lrn;
string lrn_grad;
string matmul;
@@ -670,7 +694,11 @@
string mkl_conv2d_grad_filter;
string mkl_conv2d_grad_filter_with_bias;
string mkl_conv2d_with_bias;
+ string mkl_fused_conv2d;
+ string mkl_pad_with_conv2d;
string mul;
+ string pad;
+ string pad_with_conv2d;
string quantized_avg_pool;
string quantized_conv2d;
string quantized_conv2d_with_requantize;
@@ -798,6 +826,7 @@
// Helper function to merge different nodes
Status MergeConv2DWithBiasAdd(std::unique_ptr<Graph>* g, Node* m, Node* n);
+ Status MergePadWithConv2D(std::unique_ptr<Graph>* g, Node* m, Node* n);
Status MergeConv2DBackpropFilterWithBiasAddGrad(std::unique_ptr<Graph>* g,
Node* m, Node* n);
@@ -835,6 +864,54 @@
return n;
}
+ // Find Pad or Conv2D node that can be merged with input node 'm'.
+ // If input 'm' is Pad, then check if there exists Conv2D node that can be
+ // merged with 'm'. If input 'm' is Conv2D, then check if there exists Pad
+ // node that can be merged with 'm'.
+ static Node* GetPadOrConv2D(const Node* m) {
+ DCHECK(m);
+ Node* n = nullptr;
+
+ const Node* conv_node;
+ if (m->type_string() == csinfo_.pad) {
+ // If m is Pad, then Conv2D is the output of Pad.
+ for (const Edge* e : m->out_edges()) {
+ if (!e->IsControlEdge() && e->dst()->type_string() == csinfo_.conv2d) {
+ n = e->dst();
+ conv_node = n;
+ break;
+ }
+ }
+ } else {
+ DCHECK_EQ(m->type_string(), csinfo_.conv2d);
+ // If m is conv2D, Go over all input edges
+ // and search for Pad Node.
+ for (const Edge* e : m->in_edges()) {
+ if (!e->IsControlEdge() && e->src()->type_string() == csinfo_.pad) {
+ n = e->src();
+ conv_node = m;
+ break;
+ }
+ }
+ }
+ // Check if only VALID type of padding is used
+ // or not.
+ if (n != nullptr) {
+ string padding;
+ TF_CHECK_OK(GetNodeAttr(conv_node->def(), "padding", &padding));
+ if (padding != "VALID")
+ // Then do not merge.
+ // Only VALID type of padding in conv op can be
+ // merged with Pad op.
+ n = nullptr;
+ } else {
+ VLOG(1) << "MklLayoutRewritePass: Could not find matching "
+ << "Pad and Conv2D node for merging. Input node: "
+ << m->DebugString();
+ }
+
+ return n;
+ }
// Find Conv2DBackpropFilter or BiasAddGrad node that can be merged with input
// node 'm'. If input 'm' is Conv2DBackpropFilter, then check if there exists
// BiasAddGrad node that can be merged with 'm'. If input 'm' is BiasAddGrad,
@@ -1081,6 +1158,30 @@
return do_rewrite;
}
+ // MKL-DNN's LeakyRelu(feature) = feature (if feature > 0), or
+ // feature * alpha (otherwise),
+ // while TensorFlow's LeakyRelu(feature) = max(feature, feature * alpha).
+ // These two algorithms are not consistent when alpha > 1,
+ // so we only rewrite LeakyRelu to MKL OP when alpha <= 1.
+ static bool LeakyReluRewrite(const Node* n) {
+ DCHECK(n);
+
+ float alpha;
+ bool has_attr = GetNodeAttr(n->def(), "alpha", &alpha).ok();
+ DCHECK(has_attr);
+
+ // If the alpha of LeakyRelu is less than 1, rewrite the node.
+ // Otherwise eigen node is used instead.
+ if (alpha <= 1) {
+ return true;
+ }
+ VLOG(1) << "LeakyReluRewrite: The model sets alpha is greater than 1 "
+ << "which case is not optimized by Intel MKL, thus using Eigen op"
+ << "for LeakyRelu ";
+
+ return false;
+ }
+
static bool MaxpoolGradRewrite(const Node* n) {
CHECK_NOTNULL(n);
bool do_rewrite = false;
@@ -1113,6 +1214,23 @@
return false;
}
+ static bool FusedConv2DRewrite(const Node* n) {
+ // MKL DNN currently doesn't support all fusions that grappler fuses
+ // together with Conv2D (ex. batchnorm). We rewrite _FusedConv2D only if
+ // it includes those we support.
+ DataType T;
+ if (!GetNodeAttr(n->def(), "T", &T).ok() ||
+ !mkl_op_registry::IsMklOp(csinfo_.mkl_fused_conv2d, T)) {
+ return false;
+ }
+
+ std::vector<string> fused_ops;
+ TF_CHECK_OK(GetNodeAttr(n->def(), "fused_ops", &fused_ops));
+ return (fused_ops == std::vector<string>{"BiasAdd"} ||
+ fused_ops == std::vector<string>{"Relu"} ||
+ fused_ops == std::vector<string>{"BiasAdd", "Relu"});
+ }
+
// Rewrites input node to a new node specified by its matching rewrite info.
//
// Method first searches matching rewrite info for input node and then
@@ -1274,8 +1392,17 @@
bool change_format = false);
static void CopyAttrsFusedBatchNorm(const Node* orig_node, NodeBuilder* nb,
bool change_format = false);
+ static void CopyAttrsLeakyRelu(const Node* orig_node, NodeBuilder* nb,
+ bool change_format = false);
+ static void CopyAttrsFusedConv2D(const Node* orig_node, NodeBuilder* nb,
+ bool change_format = false);
static void CopyAttrsLRN(const Node* orig_node, NodeBuilder* nb,
bool change_format = false);
+ static void CopyAttrsPadWithConv2D(const Node* orig_node, NodeBuilder* nb,
+ bool change_format = false);
+ static void CopyAttrsFromPadAndConv2D(const Node* orig_node1,
+ const Node* orig_node2, NodeBuilder* nb,
+ bool change_format = false);
static void CopyAttrsPooling(const Node* orig_node, NodeBuilder* nb,
bool change_format = false);
static void CopyAttrsQuantizedPooling(const Node* orig_node, NodeBuilder* nb,
@@ -1488,10 +1615,13 @@
CHECK_NOTNULL(filter_node);
// Now check which nodes receive from filter_node. Filter feeds as
- // 2nd input (slot 1) of _MklConv2D and _MklConv2DWithBias.
+ // 2nd input (slot 1) of _MklConv2D, _MklConv2DWithBias, and
+ // _MklFusedConv2D.
for (const Edge* e : filter_node->out_edges()) {
if ((e->dst()->type_string() == csinfo_.mkl_conv2d ||
- e->dst()->type_string() == csinfo_.mkl_conv2d_with_bias) &&
+ e->dst()->type_string() == csinfo_.mkl_pad_with_conv2d ||
+ e->dst()->type_string() == csinfo_.mkl_conv2d_with_bias ||
+ e->dst()->type_string() == csinfo_.mkl_fused_conv2d) &&
e->dst_input() == kConv2DFilterInputSlotIdx
/* filter is 2nd input of Conv2D and _MklConv2D. */) {
if (conv2d_node != nullptr) {
@@ -1846,6 +1976,72 @@
}
}
+// Used in rinfo when replacing __MklDummyPadWithConv2D by _MklPadWithConv2D
+void MklLayoutRewritePass::CopyAttrsPadWithConv2D(const Node* orig_node,
+ NodeBuilder* nb,
+ bool change_format) {
+ DataType Tpaddings;
+ DataType T;
+ string data_format;
+ string padding;
+ std::vector<int32> strides;
+ std::vector<int32> dilations;
+ bool use_cudnn_on_gpu;
+
+ // Get all attributes from old node.
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "T", &T));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "strides", &strides));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "dilations", &dilations));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "padding", &padding));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "data_format", &data_format));
+ TF_CHECK_OK(
+ GetNodeAttr(orig_node->def(), "use_cudnn_on_gpu", &use_cudnn_on_gpu));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "Tpaddings", &Tpaddings));
+
+ // Add attributes to new node.
+ nb->Attr("T", T);
+ nb->Attr("strides", strides);
+ nb->Attr("dilations", dilations);
+ nb->Attr("padding", padding);
+ nb->Attr("data_format", data_format);
+ nb->Attr("use_cudnn_on_gpu", use_cudnn_on_gpu);
+ nb->Attr("Tpaddings", Tpaddings);
+}
+
+// Used with MergePadWithConv2D
+void MklLayoutRewritePass::CopyAttrsFromPadAndConv2D(const Node* orig_node1,
+ const Node* orig_node2,
+ NodeBuilder* nb,
+ bool change_format) {
+ DataType Tpaddings;
+ DataType T;
+ string data_format;
+ string padding;
+ std::vector<int32> strides;
+ std::vector<int32> dilations;
+ bool use_cudnn_on_gpu;
+
+ // Get all attributes from old node 1.
+ TF_CHECK_OK(GetNodeAttr(orig_node1->def(), "T", &T));
+ TF_CHECK_OK(GetNodeAttr(orig_node1->def(), "strides", &strides));
+ TF_CHECK_OK(GetNodeAttr(orig_node1->def(), "dilations", &dilations));
+ TF_CHECK_OK(GetNodeAttr(orig_node1->def(), "padding", &padding));
+ TF_CHECK_OK(GetNodeAttr(orig_node1->def(), "data_format", &data_format));
+ TF_CHECK_OK(
+ GetNodeAttr(orig_node1->def(), "use_cudnn_on_gpu", &use_cudnn_on_gpu));
+ // Get all attributes from old node 2.
+ TF_CHECK_OK(GetNodeAttr(orig_node2->def(), "Tpaddings", &Tpaddings));
+
+ // Add attributes to new node.
+ nb->Attr("T", T);
+ nb->Attr("strides", strides);
+ nb->Attr("dilations", dilations);
+ nb->Attr("padding", padding);
+ nb->Attr("data_format", data_format);
+ nb->Attr("use_cudnn_on_gpu", use_cudnn_on_gpu);
+ nb->Attr("Tpaddings", Tpaddings);
+}
+
void MklLayoutRewritePass::CopyAttrsAddN(const Node* orig_node, NodeBuilder* nb,
bool change_format) {
DataType T;
@@ -1901,6 +2097,21 @@
nb->Attr("beta", beta);
}
+void MklLayoutRewritePass::CopyAttrsLeakyRelu(const Node* orig_node,
+ NodeBuilder* nb,
+ bool change_format) {
+ DataType T;
+ float alpha;
+
+ // Get all attributes from old node.
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "T", &T));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "alpha", &alpha));
+
+ // Add attributes to new node.
+ nb->Attr("T", T);
+ nb->Attr("alpha", alpha);
+}
+
void MklLayoutRewritePass::CopyAttrsPooling(const Node* orig_node,
NodeBuilder* nb,
bool change_format) {
@@ -2100,6 +2311,39 @@
nb->Attr("is_training", is_training);
}
+void MklLayoutRewritePass::CopyAttrsFusedConv2D(const Node* orig_node,
+ NodeBuilder* nb,
+ bool change_format) {
+ DataType T;
+ int num_args;
+ float epsilon;
+ string data_format;
+ string padding;
+ std::vector<int32> strides;
+ std::vector<int32> dilations;
+ std::vector<string> fused_ops;
+
+ // Get all attributes from old node.
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "T", &T));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "num_args", &num_args));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "strides", &strides));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "padding", &padding));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "data_format", &data_format));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "dilations", &dilations));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "fused_ops", &fused_ops));
+ TF_CHECK_OK(GetNodeAttr(orig_node->def(), "epsilon", &epsilon));
+
+ // Add attributes to new node.
+ nb->Attr("T", T);
+ nb->Attr("num_args", num_args);
+ nb->Attr("strides", strides);
+ nb->Attr("padding", padding);
+ nb->Attr("data_format", data_format);
+ nb->Attr("dilations", dilations);
+ nb->Attr("fused_ops", fused_ops);
+ nb->Attr("epsilon", epsilon);
+}
+
//////////////////////////////////////////////////////////////////////////
// Helper functions related to node merge pass
//////////////////////////////////////////////////////////////////////////
@@ -2297,6 +2541,165 @@
return Status::OK();
}
+Status MklLayoutRewritePass::MergePadWithConv2D(std::unique_ptr<Graph>* g,
+ Node* m, Node* n) {
+ DCHECK(((m->type_string() == csinfo_.pad &&
+ n->type_string() == csinfo_.conv2d)) ||
+ ((n->type_string() == csinfo_.pad &&
+ m->type_string() == csinfo_.conv2d)));
+
+ // Conv2D is successor node, and Pad predecessor node.
+ Node* pred = m->type_string() == csinfo_.pad ? m : n;
+ Node* succ = m->type_string() == csinfo_.pad ? n : m;
+
+ // 1. Get all attributes from input nodes.
+ DataType T_pred, T_succ;
+ string padding;
+ std::vector<int32> strides;
+ std::vector<int32> dilations;
+ string data_format_pred, data_format_succ;
+ bool use_cudnn_on_gnu;
+ TF_CHECK_OK(GetNodeAttr(pred->def(), "T", &T_pred));
+ TF_CHECK_OK(GetNodeAttr(succ->def(), "T", &T_succ));
+ TF_CHECK_OK(GetNodeAttr(succ->def(), "padding", &padding));
+ TF_CHECK_OK(GetNodeAttr(succ->def(), "strides", &strides));
+ TF_CHECK_OK(GetNodeAttr(succ->def(), "dilations", &dilations));
+ // Data format for pad is not available and not necessary, thus
+ // dont need to match data format for Pad
+ TF_CHECK_OK(GetNodeAttr(succ->def(), "data_format", &data_format_succ));
+ TF_CHECK_OK(GetNodeAttr(succ->def(), "use_cudnn_on_gpu", &use_cudnn_on_gnu));
+ // Check if the data types and devices of both succ and pred are the same.
+ // Assert is not used, because it can be too strict.
+ // Don't need to check for data formats because it is not available in Pad.
+ if (T_pred != T_succ ||
+ pred->assigned_device_name() != succ->assigned_device_name() ||
+ pred->def().device() != succ->def().device()) {
+ return Status(error::Code::INVALID_ARGUMENT,
+ "T attribute or devices of Conv2D and "
+ "Pad do not match. Will skip node merge optimization");
+ }
+
+ const int succ_num = succ->num_inputs();
+ gtl::InlinedVector<Node*, 4> succ_control_edges;
+ gtl::InlinedVector<std::pair<Node*, int>, 4> succ_in(succ_num);
+ FillInputs(succ, &succ_control_edges, &succ_in);
+
+ const int pred_num = pred->num_inputs();
+ gtl::InlinedVector<Node*, 4> pred_control_edges;
+ gtl::InlinedVector<std::pair<Node*, int>, 4> pred_in(pred_num);
+ FillInputs(pred, &pred_control_edges, &pred_in);
+
+ // We need to ensure that Pad only feeds to Conv2D (some other operator is
+ // not expecting output of Pad). If this is not the case, then we cannot
+ // merge Conv2D with Pad.
+ const int kFirstOutputSlot = 0;
+ for (const Edge* e : pred->out_edges()) {
+ if (e->src_output() == kFirstOutputSlot && e->dst() != succ) {
+ return Status(error::Code::INVALID_ARGUMENT,
+ "Pad does not feed to Conv2D, or "
+ "it feeds Conv2D but has multiple outputs. "
+ "Will skip node merge optimization");
+ }
+ }
+
+ // 2. Get inputs from both the nodes.
+
+ // Pad must have 2 data inputs: "input" and paddings.
+ int PadDataInputEdges = 0;
+ for (const Edge* e : pred->in_edges()) {
+ if (!e->IsControlEdge()) {
+ PadDataInputEdges++;
+ }
+ }
+ DCHECK_EQ(PadDataInputEdges, 2);
+
+ // Conv2D must have 2 data inputs: pad output and Filter
+ int ConvDataInputEdges = 0;
+ for (const Edge* e : succ->in_edges()) {
+ if (!e->IsControlEdge()) {
+ ConvDataInputEdges++;
+ }
+ }
+ DCHECK_EQ(ConvDataInputEdges, 2);
+
+ // We will use the node name of Conv2D as the name of new node
+ // Build new node. We use same name as original node, but change the op
+ // name.
+ NodeBuilder nb(succ->name(), csinfo_.pad_with_conv2d);
+ nb.Input(pred_in[0].first, pred_in[0].second); // In1 (input data) of Pad
+ // pred_in[1] will be 2nd Tensorflow tensor for Conv2D.
+ nb.Input(succ_in[1].first, succ_in[1].second); // In2 (filter) of conv2d
+ // In1 of Conv2D is same as output of Pad.
+ // Thus, only need to add In2 of Conv2D
+ nb.Input(pred_in[1].first, pred_in[1].second); // In2 (paddings) of Pad
+
+ // Copy attributes from Pad and conv2D to PadWithConv2D.
+ CopyAttrsFromPadAndConv2D(const_cast<const Node*>(succ),
+ const_cast<const Node*>(pred), &nb);
+
+ // Copy the device assigned to old node to new node.
+ nb.Device(succ->def().device());
+
+ // Create node.
+ Node* new_node;
+ TF_CHECK_OK(nb.Finalize(&**g, &new_node));
+ DCHECK(new_node);
+
+ // Incoming data edges from 'pred' node and 'succ' node to new 'new_node'
+ // node are already copied in BuildNode.
+ // We handle control edges now.
+ for (const Edge* e : pred->in_edges()) {
+ if (e->IsControlEdge()) {
+ // Don't allow duplicate edge
+ (*g)->AddControlEdge(e->src(), new_node, false);
+ }
+ }
+ for (const Edge* e : succ->in_edges()) {
+ if (e->IsControlEdge()) {
+ // Don't allow duplicate edge
+ (*g)->AddControlEdge(e->src(), new_node, false);
+ }
+ }
+
+ // Incoming edges are fixed, we will fix the outgoing edges now.
+ // First, we will fix outgoing control edges from 'pred' node.
+ for (const Edge* e : pred->out_edges()) {
+ if (e->IsControlEdge()) {
+ // Don't allow duplicate edge
+ (*g)->AddControlEdge(new_node, e->dst(), false);
+ }
+ }
+
+ // Second, we will fix outgoing control and data edges from 'succ' node.
+ for (const Edge* e : succ->out_edges()) {
+ if (e->IsControlEdge()) {
+ // Allow duplicate while adding control edge as it would fail (return
+ // NULL) if we try to add duplicate edge.
+ (*g)->AddControlEdge(new_node, e->dst(), false);
+ } else {
+ // Conv2D has only 1 output (at slot 0) and merged node also has only 1
+ // output (at slot 0).
+ const int kPadWithConv2DOutputSlot = 0;
+ (*g)->AddEdge(new_node, kPadWithConv2DOutputSlot, e->dst(),
+ e->dst_input());
+ }
+ }
+
+ // Copy device assigned to old node to new node.
+ // It's ok to use pred or succ as we have enforced a check that
+ // both have same device assigned.
+ new_node->set_assigned_device_name(pred->assigned_device_name());
+
+ VLOG(1) << "MklLayoutRewritePass: Merged old node:" << pred->DebugString()
+ << ", and node: " << succ->DebugString()
+ << ", into node:" << new_node->DebugString();
+
+ (*g)->RemoveNode(succ);
+ (*g)->RemoveNode(pred);
+
+ return Status::OK();
+}
+
Status MklLayoutRewritePass::MergeConv2DBackpropFilterWithBiasAddGrad(
std::unique_ptr<Graph>* g, Node* m, Node* n) {
CHECK_EQ(((m->type_string() == csinfo_.bias_add_grad &&
@@ -2430,6 +2833,12 @@
m->type_string() == csinfo_.conv2d))) {
return this->MergeConv2DWithBiasAdd(g, m, n);
}
+ if (((m->type_string() == csinfo_.pad &&
+ n->type_string() == csinfo_.conv2d)) ||
+ ((n->type_string() == csinfo_.pad &&
+ m->type_string() == csinfo_.conv2d))) {
+ return this->MergePadWithConv2D(g, m, n);
+ }
if (((m->type_string() == csinfo_.bias_add_grad &&
n->type_string() == csinfo_.conv2d_grad_filter)) ||
@@ -2576,11 +2985,13 @@
return nullptr;
}
- // We make an exception for __MklDummyConv2DWithBias and
- // __MklConv2DBackpropFilterWithBias since their names do not match Mkl node
- // names.
+ // We make an exception for __MklDummyConv2DWithBias,
+ // __MklConv2DBackpropFilterWithBias, and __MklDummyPadWithConv2D since their
+ // names do not match Mkl node names.
if (n->type_string() != csinfo_.conv2d_with_bias &&
+ n->type_string() != csinfo_.pad_with_conv2d &&
n->type_string() != csinfo_.conv2d_grad_filter_with_bias &&
+ n->type_string() != csinfo_.fused_conv2d &&
!mkl_op_registry::IsMklOp(mkl_op_registry::GetMklOpName(n->type_string()),
T)) {
return nullptr;
diff --git a/tensorflow/core/graph/mkl_layout_pass_test.cc b/tensorflow/core/graph/mkl_layout_pass_test.cc
index 22f8791..197ec0c 100644
--- a/tensorflow/core/graph/mkl_layout_pass_test.cc
+++ b/tensorflow/core/graph/mkl_layout_pass_test.cc
@@ -65,6 +65,13 @@
class MklLayoutPassTest : public ::testing::Test {
public:
MklLayoutPassTest() : graph_(OpRegistry::Global()) {}
+ // Ashraf added
+ Node* FindNode(const string& name) {
+ for (Node* node : graph_.nodes()) {
+ if (node->name() == name) return node;
+ }
+ LOG(FATAL) << name;
+ }
void InitGraph(const string& s, const string& device = kCPUDevice) {
::tensorflow::InitGraph(s, &graph_, device);
@@ -126,14 +133,17 @@
REGISTER_OP("InputList").Output("o: N * float").Attr("N: int").SetIsStateful();
REGISTER_OP("HalfInput").Output("o: half").SetIsStateful();
REGISTER_OP("Int32Input").Output("o: int32").SetIsStateful();
+REGISTER_OP("DoubleInput").Output("o: double").SetIsStateful();
REGISTER_OP("_MklInput").Output("o: uint8").SetIsStateful();
REGISTER_OP("_MklInput2")
.Output("o: uint8")
.Output("o1: uint8")
.SetIsStateful();
+REGISTER_OP("Output2").Input("i: float").Input("i1: float").SetIsStateful();
+REGISTER_OP("Output").Input("i: float").SetIsStateful();
/////////////////////////////////////////////////////////////////////
-// Unit tests related to node merge optiimization
+// Unit tests related to node merge optimization
/////////////////////////////////////////////////////////////////////
TEST_F(MklLayoutPassTest, Basic) {
@@ -455,6 +465,248 @@
"E:3->G:4;F->G;F:control->DMT/_3:control;G->Z;X->Y:1;X->Z:1");
}
+// Test set 3: Pad + Conv2D fusion
+// padding is VALID type
+// A = input(image), B = input(paddings), C= Pad = input of conv2D,
+// D=input(filter), E = Conv2D, Z = Zeta
+// C=Pad(A,B); E=Conv2D(C,D); Z=Zeta(E,Y)
+// After layout pass
+// _MklPadWithConv2D(A, D, B, DMT/_0, DMT/_1, DMT/_2)
+TEST_F(MklLayoutPassTest, NodeMerge_PadWithConv2D_Positive) {
+ DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Pad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Tpaddings' value { type: DT_INT32 } }"
+ " input: ['A', 'B']}"
+ "node { name: 'D' op: 'Input'}"
+ "node { name: 'E' op: 'Conv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'data_format' value { s: 'NHWC' } }"
+ " attr { key: 'use_cudnn_on_gpu' value { b: false } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'VALID' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " input: ['C', 'D'] }"
+ "node { name: 'Y' op: 'Input'}"
+ "node { name: 'Z' op: 'Zeta'"
+ " attr {key: 'T' value { type: DT_FLOAT } }"
+ " input: ['E', 'Y']}");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Int32Input);D(Input);DMT/_0(Const);DMT/_1(Const);"
+ "DMT/_2(Const);E(_MklPadWithConv2D);Y(Input);Z(Zeta)|A->E;"
+ "A:control->DMT/_0:control;A:control->DMT/_1:control;"
+ "A:control->DMT/_2:control;B->E:2;D->E:1;DMT/_0->E:3;DMT/_1->E:4;"
+ "DMT/_2->E:5;E->Z;Y->Z:1");
+}
+// Test if input control edges do not duplicate after merge.
+// If both the merging ops have input control edge from a common op
+// then, the merged op will have only one control edge from that
+// common op.
+// padding is VALID type
+// A = input(image), A1 = input, B = input(paddings),
+// C= Pad = input of conv2D,
+// D=input(filter), E = Conv2D, Z = Zeta
+// C=Pad(A,B); E=Conv2D(C,D); Z=Zeta(E,Y)
+// A1:control->C:control
+// A1:control->E:control
+// After layout pass:
+// _MklPadWithConv2D(A, D, B, DMT/_0, DMT/_1, DMT/_2)
+// A1:control->E:control (only one control edge)
+TEST_F(MklLayoutPassTest, Input_ControlEdge_PadWithConv2D_Positive) {
+ DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
+ InitGraph(
+ "node { name: 'A1' op: 'Input'}"
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Pad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Tpaddings' value { type: DT_INT32 } }"
+ " input: ['A', 'B']}"
+ "node { name: 'D' op: 'Input'}"
+ "node { name: 'E' op: 'Conv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'data_format' value { s: 'NHWC' } }"
+ " attr { key: 'use_cudnn_on_gpu' value { b: false } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'VALID' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " input: ['C', 'D'] }"
+ "node { name: 'Y' op: 'Input'}"
+ "node { name: 'Z' op: 'Zeta'"
+ " attr {key: 'T' value { type: DT_FLOAT } }"
+ " input: ['E', 'Y']}");
+ Node* a1 = FindNode("A1");
+ Node* c = FindNode("C");
+ Node* e = FindNode("E");
+ const Edge* edge = graph_.AddControlEdge(a1, c);
+ const Edge* edge_1 = graph_.AddControlEdge(a1, e);
+ ASSERT_NE(edge, nullptr);
+ ASSERT_NE(edge_1, nullptr);
+ EXPECT_EQ(
+ DoMklLayoutOptimizationPass(),
+ "A(Input);A1(Input);B(Int32Input);D(Input);DMT/_0(Const);DMT/_1(Const);"
+ "DMT/_2(Const);E(_MklPadWithConv2D);Y(Input);Z(Zeta)|A->E;"
+ "A1:control->E:control;A:control->DMT/_0:control;A:control->DMT/"
+ "_1:control;"
+ "A:control->DMT/_2:control;B->E:2;D->E:1;DMT/_0->E:3;DMT/_1->E:4;"
+ "DMT/_2->E:5;E->Z;Y->Z:1");
+}
+// Test if output control edges does not duplicate after merge.
+// If both the merging ops have output control edge to a common op,
+// then after merge, the merged op will have only one control edge
+// to that commom op.
+// padding is VALID type
+// A = input(image), B = input(paddings), C= Pad = input of conv2D,
+// D=input(filter), E = Conv2D, Z = Zeta
+// C=Pad(A,B); E=Conv2D(C,D); Z=Zeta(E,Y)
+// C:control->A1:control
+// E:control->A1:control
+// After layout pass:
+// _MklPadWithConv2D(A, D, B, DMT/_0, DMT/_1, DMT/_2)
+// E:control->A1:control (only one control edge)
+TEST_F(MklLayoutPassTest, Output_ControlEdge_PadWithConv2D_Positive) {
+ DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
+ InitGraph(
+ "node { name: 'A1' op: 'Input'}"
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Pad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Tpaddings' value { type: DT_INT32 } }"
+ " input: ['A', 'B']}"
+ "node { name: 'D' op: 'Input'}"
+ "node { name: 'E' op: 'Conv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'data_format' value { s: 'NHWC' } }"
+ " attr { key: 'use_cudnn_on_gpu' value { b: false } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'VALID' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " input: ['C', 'D'] }"
+ "node { name: 'Y' op: 'Input'}"
+ "node { name: 'Z' op: 'Zeta'"
+ " attr {key: 'T' value { type: DT_FLOAT } }"
+ " input: ['E', 'Y']}");
+ Node* a1 = FindNode("A1");
+ Node* c = FindNode("C");
+ Node* e = FindNode("E");
+ const Edge* edge = graph_.AddControlEdge(c, a1);
+ const Edge* edge_1 = graph_.AddControlEdge(e, a1);
+ ASSERT_NE(edge, nullptr);
+ ASSERT_NE(edge_1, nullptr);
+ EXPECT_EQ(
+ DoMklLayoutOptimizationPass(),
+ "A(Input);A1(Input);B(Int32Input);D(Input);DMT/_0(Const);DMT/_1(Const);"
+ "DMT/_2(Const);E(_MklPadWithConv2D);Y(Input);Z(Zeta)|A->E;"
+ "A:control->DMT/_0:control;A:control->DMT/_1:control;"
+ "A:control->DMT/_2:control;B->E:2;D->E:1;DMT/_0->E:3;DMT/_1->E:4;"
+ "DMT/_2->E:5;E->Z;E:control->A1:control;Y->Z:1");
+}
+// Pad + Conv2D fusion with padding is VALID,
+// Input node pointing to both Pad and Conv2D
+// A = input(image), B = input(paddings), C= Pad
+// E = Conv2D, Z = Zeta
+// C=Pad(A,B); E=Conv2D(C,A); Z=Zeta(E,Y)
+// After layout pass
+// _MklPadWithConv2D(A, A, B, DMT/_0, DMT/_1, DMT/_2)
+TEST_F(MklLayoutPassTest, NodeMerge_PadWithConv2D_Common_Input) {
+ DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Pad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Tpaddings' value { type: DT_INT32 } }"
+ " input: ['A', 'B']}"
+ "node { name: 'E' op: 'Conv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'data_format' value { s: 'NHWC' } }"
+ " attr { key: 'use_cudnn_on_gpu' value { b: false } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'VALID' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " input: ['C', 'A'] }"
+ "node { name: 'Y' op: 'Input'}"
+ "node { name: 'Z' op: 'Zeta'"
+ " attr {key: 'T' value { type: DT_FLOAT } }"
+ " input: ['E', 'Y']}");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Int32Input);DMT/_0(Const);DMT/_1(Const);"
+ "DMT/_2(Const);E(_MklPadWithConv2D);Y(Input);Z(Zeta)|A->E;A->E:1;"
+ "A:control->DMT/_0:control;A:control->DMT/_1:control;"
+ "A:control->DMT/_2:control;B->E:2;DMT/_0->E:3;DMT/_1->E:4;"
+ "DMT/_2->E:5;E->Z;Y->Z:1");
+}
+// Pad + Conv2D with padding is VALID,
+// Input node pointing to both Pad and Conv2D
+// Output of both Pad and Conv2D feeds one node (Z as Output2)
+// A = input(as image), B = input(as paddings), C= Pad
+// E = Conv2D, Z = Output2
+// C=Pad(A,B); E=Conv2D(C,A); Z=Output(C,E)
+// After layout pass - No merging, since Pad and Conv2D both
+// feed to the same node (Z)
+TEST_F(MklLayoutPassTest, NodeMerge_PadWithConv2D_Common_InOutput) {
+ DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Pad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Tpaddings' value { type: DT_INT32 } }"
+ " input: ['A', 'B']}"
+ "node { name: 'E' op: 'Conv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'data_format' value { s: 'NHWC' } }"
+ " attr { key: 'use_cudnn_on_gpu' value { b: false } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'VALID' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " input: ['C', 'A'] }"
+ "node { name: 'Z' op: 'Output2'"
+ " input: ['C', 'E']}");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Int32Input);C(Pad);DMT/_0(Const);DMT/_1(Const);"
+ "E(_MklConv2D);Z(Output2)|A->C;A->E:1;B->C:1;C->E;C->Z;"
+ "C:control->DMT/_0:control;C:control->DMT/_1:control;"
+ "DMT/_0->E:2;DMT/_1->E:3;E->Z:1");
+}
+// Pad + Conv2D; padding is SAME
+// A = input(image), B = input(paddings), C= Pad = input of conv2D,
+// D=input(filter), E = Conv2D, Z = Zeta
+// C=Pad(A,B); E=Conv2D(C,D); Z=Zeta(E,Y)
+// After layout pass - No merging
+TEST_F(MklLayoutPassTest, NodeMerge_PadWithConv2D_Negative) {
+ DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Int32Input'}"
+ "node { name: 'C' op: 'Pad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'Tpaddings' value { type: DT_INT32 } }"
+ " input: ['A', 'B']}"
+ "node { name: 'D' op: 'Input'}"
+ "node { name: 'E' op: 'Conv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'data_format' value { s: 'NHWC' } }"
+ " attr { key: 'use_cudnn_on_gpu' value { b: false } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'SAME' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " input: ['C', 'D'] }"
+ "node { name: 'Y' op: 'Input'}"
+ "node { name: 'Z' op: 'Zeta'"
+ " attr {key: 'T' value { type: DT_FLOAT } }"
+ " input: ['E', 'Y']}");
+ EXPECT_EQ(
+ DoMklLayoutOptimizationPass(),
+ "A(Input);B(Int32Input);C(Pad);D(Input);DMT/_0(Const);DMT/_1(Const);"
+ "E(_MklConv2D);Y(Input);Z(Zeta)|A->C;B->C:1;C->E;"
+ "C:control->DMT/_0:control;C:control->DMT/_1:control;"
+ "D->E:1;DMT/_0->E:2;DMT/_1->E:3;E->Z;Y->Z:1");
+}
#ifdef ENABLE_TRANSPOSE_OPTIMIZATION
TEST_F(MklLayoutPassTest, NodeMerge_TransposeConv2DTranspose_Positive) {
InitGraph(
@@ -845,6 +1097,131 @@
"A->C;B->C:1;B->D;C->D:1");
}
+// Rewrite test for _FusedConv2D Op with BiasAdd fusion
+TEST_F(MklLayoutPassTest, NodeRewrite_FusedConv2D_Positive1) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Input'}"
+ "node { name: 'C' op: 'Input'}"
+ "node { name: 'D' op: '_FusedConv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'num_args' value { i: 1 } }"
+ " attr { key: 'data_format' value { s: 'NCHW' } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'SAME' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'fused_ops' value { list: {s: 'BiasAdd'} } }"
+ " attr { key: 'epsilon' value { f: 0.001 }}"
+ " input: ['A', 'B', 'C']}"
+ "node { name: 'E' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['D', 'C'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Input);C(Input);D(_MklFusedConv2D);DMT/_0(Const);"
+ "DMT/_1(Const);DMT/_2(Const);E(Zeta)|A->D;"
+ "A:control->DMT/_0:control;A:control->DMT/_1:control;"
+ "A:control->DMT/_2:control;B->D:1;C->D:2;C->E:1;D->E;"
+ "DMT/_0->D:3;DMT/_1->D:4;DMT/_2->D:5");
+}
+
+// Rewrite test for _FusedConv2D Op with Relu fusion
+TEST_F(MklLayoutPassTest, NodeRewrite_FusedConv2D_Positive2) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Input'}"
+ "node { name: 'C' op: 'Input'}"
+ "node { name: 'D' op: '_FusedConv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'num_args' value { i: 1 } }"
+ " attr { key: 'data_format' value { s: 'NCHW' } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'SAME' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'fused_ops' value { list: {s: 'Relu'} } }"
+ " attr { key: 'epsilon' value { f: 0.001 }}"
+ " input: ['A', 'B', 'C']}"
+ "node { name: 'E' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['D', 'C'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Input);C(Input);D(_MklFusedConv2D);DMT/_0(Const);"
+ "DMT/_1(Const);DMT/_2(Const);E(Zeta)|A->D;"
+ "A:control->DMT/_0:control;A:control->DMT/_1:control;"
+ "A:control->DMT/_2:control;B->D:1;C->D:2;C->E:1;D->E;"
+ "DMT/_0->D:3;DMT/_1->D:4;DMT/_2->D:5");
+}
+
+// Rewrite test for _FusedConv2D Op with BiasAdd+Relu fusion
+TEST_F(MklLayoutPassTest, NodeRewrite_FusedConv2D_Positive3) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Input'}"
+ "node { name: 'C' op: 'Input'}"
+ "node { name: 'D' op: '_FusedConv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'num_args' value { i: 1 } }"
+ " attr { key: 'data_format' value { s: 'NCHW' } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'SAME' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'fused_ops'"
+ " value { list: {s: 'BiasAdd', s: 'Relu'} } }"
+ " attr { key: 'epsilon' value { f: 0.001 }}"
+ " input: ['A', 'B', 'C']}"
+ "node { name: 'E' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['D', 'C'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Input);C(Input);D(_MklFusedConv2D);DMT/_0(Const);"
+ "DMT/_1(Const);DMT/_2(Const);E(Zeta)|A->D;"
+ "A:control->DMT/_0:control;A:control->DMT/_1:control;"
+ "A:control->DMT/_2:control;B->D:1;C->D:2;C->E:1;D->E;"
+ "DMT/_0->D:3;DMT/_1->D:4;DMT/_2->D:5");
+}
+
+// Rewrite test for _FusedConv2D Op with unsupported fusion
+TEST_F(MklLayoutPassTest, NodeRewrite_FusedConv2D_Negative1) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Input'}"
+ "node { name: 'C' op: 'Input'}"
+ "node { name: 'D' op: '_FusedConv2D'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'num_args' value { i: 1 } }"
+ " attr { key: 'data_format' value { s: 'NCHW' } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'SAME' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'fused_ops' value { list: {s: 'Unsupported'} } }"
+ " attr { key: 'epsilon' value { f: 0.001 }}"
+ " input: ['A', 'B', 'C']}"
+ "node { name: 'E' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['D', 'C'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Input);C(Input);D(_FusedConv2D);E(Zeta)|A->D;"
+ "B->D:1;C->D:2;C->E:1;D->E");
+}
+
+// Rewrite test for _FusedConv2D Op with unsupported type
+TEST_F(MklLayoutPassTest, NodeRewrite_FusedConv2D_Negative2) {
+ InitGraph(
+ "node { name: 'A' op: 'DoubleInput'}"
+ "node { name: 'B' op: 'DoubleInput'}"
+ "node { name: 'C' op: 'DoubleInput'}"
+ "node { name: 'D' op: '_FusedConv2D'"
+ " attr { key: 'T' value { type: DT_DOUBLE } }"
+ " attr { key: 'num_args' value { i: 1 } }"
+ " attr { key: 'data_format' value { s: 'NCHW' } }"
+ " attr { key: 'strides' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'padding' value { s: 'SAME' } }"
+ " attr { key: 'dilations' value { list: {i: 1, i:1, i:1, i:1} } }"
+ " attr { key: 'fused_ops' value { list: {s: 'BiasAdd'} } }"
+ " attr { key: 'epsilon' value { f: 0.001 }}"
+ " input: ['A', 'B', 'C']}"
+ "node { name: 'E' op: 'Zeta' attr { key: 'T' value { type: DT_DOUBLE } }"
+ " input: ['D', 'C'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(DoubleInput);B(DoubleInput);C(DoubleInput);"
+ "D(_FusedConv2D);E(Zeta)|A->D;B->D:1;C->D:2;C->E:1;D->E");
+}
+
TEST_F(MklLayoutPassTest, NodeRewrite_Conv2DGradFilter_Positive) {
InitGraph(
"node { name: 'A' op: 'Input'}"
@@ -1271,6 +1648,85 @@
"DMT/_1->C:2");
}
+TEST_F(MklLayoutPassTest, NodeRewrite_LeakyRelu_Positive) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'LeakyRelu'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'alpha' value { f: 0.1 } }"
+ " input: ['A'] }"
+ "node { name: 'C' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['A', 'B'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(_MklLeakyRelu);C(Zeta);DMT/_0(Const)|A->B;A->C;"
+ "A:control->DMT/_0:control;B->C:1;DMT/_0->B:1");
+}
+
+TEST_F(MklLayoutPassTest, NodeRewrite_LeakyRelu_Negative) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'LeakyRelu'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'alpha' value { f: 2.0 } }"
+ " input: ['A'] }"
+ "node { name: 'C' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['A', 'B'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(LeakyRelu);C(Zeta)|A->B;A->C;B->C:1");
+}
+
+TEST_F(MklLayoutPassTest, NodeRewrite_LeakyReluGrad_Positive) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Input'}"
+ "node { name: 'C' op: 'LeakyReluGrad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'alpha' value { f: 0.1 } }"
+ " input: ['A', 'B'] }"
+ "node { name: 'D' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['A', 'C'] }");
+ EXPECT_EQ(DoMklLayoutOptimizationPass(),
+ "A(Input);B(Input);C(_MklLeakyReluGrad);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");
+}
+
+TEST_F(MklLayoutPassTest, NodeRewrite_LeakyReluGrad_Negative) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'Input'}"
+ "node { name: 'C' op: 'LeakyReluGrad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'alpha' value { f: 2.0 } }"
+ " input: ['A', 'B'] }"
+ "node { name: 'D' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['A', 'C'] }");
+ EXPECT_EQ(
+ DoMklLayoutOptimizationPass(),
+ "A(Input);B(Input);C(LeakyReluGrad);D(Zeta)|A->C;A->D;B->C:1;C->D:1");
+}
+
+TEST_F(MklLayoutPassTest, NodeRewrite_LeakyReluLeakyReluGrad_Positive) {
+ InitGraph(
+ "node { name: 'A' op: 'Input'}"
+ "node { name: 'B' op: 'LeakyRelu'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'alpha' value { f: 0.1 } }"
+ " input: ['A'] }"
+ "node { name: 'C' op: 'LeakyReluGrad'"
+ " attr { key: 'T' value { type: DT_FLOAT } }"
+ " attr { key: 'alpha' value { f: 0.1 } }"
+ " input: ['A', 'B'] }"
+ "node { name: 'D' op: 'Zeta' attr { key: 'T' value { type: DT_FLOAT } }"
+ " input: ['A', 'C'] }");
+ EXPECT_EQ(
+ DoMklLayoutOptimizationPass(),
+ "A(Input);B(_MklLeakyRelu);C(_MklLeakyReluGrad);D(Zeta);DMT/_0(Const);"
+ "DMT/_1(Const)|A->B;A->C;A->D;A:control->DMT/_0:control;"
+ "A:control->DMT/_1:control;B->C:1;B:1->C:3;C->D:1;DMT/_0->B:1;"
+ "DMT/_1->C:2");
+}
+
TEST_F(MklLayoutPassTest, NodeRewrite_AvgPool_Positive) {
InitGraph(
"node { name: 'A' op: 'Input'}"
diff --git a/tensorflow/core/grappler/costs/BUILD b/tensorflow/core/grappler/costs/BUILD
index 5090e62..f8af123 100644
--- a/tensorflow/core/grappler/costs/BUILD
+++ b/tensorflow/core/grappler/costs/BUILD
@@ -132,9 +132,6 @@
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
- cuda_deps = [
- "@local_config_cuda//cuda:cudnn_header",
- ],
visibility = ["//visibility:public"],
deps = [
"//third_party/eigen3",
diff --git a/tensorflow/core/grappler/costs/utils.cc b/tensorflow/core/grappler/costs/utils.cc
index 7d868a3..d45bb14 100644
--- a/tensorflow/core/grappler/costs/utils.cc
+++ b/tensorflow/core/grappler/costs/utils.cc
@@ -20,12 +20,6 @@
#include "third_party/eigen3/Eigen/Core"
-#if GOOGLE_CUDA
-#include "cuda/include/cuda.h"
-#include "cuda/include/cuda_runtime_api.h"
-#include "cuda/include/cudnn.h"
-#endif
-
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
diff --git a/tensorflow/core/grappler/graph_view.h b/tensorflow/core/grappler/graph_view.h
index 0a47b22..16156d0 100644
--- a/tensorflow/core/grappler/graph_view.h
+++ b/tensorflow/core/grappler/graph_view.h
@@ -111,32 +111,37 @@
GraphDefT* graph() const { return graph_; }
- // Find a node by name or return `nullptr` if it's not in a graph view.
+ // Finds a node by name or return `nullptr` if it's not in the graph view.
NodeDefT* GetNode(absl::string_view node_name) const {
return gtl::FindWithDefault(nodes_, node_name, nullptr);
}
- // Get the specified input port. Note that the special '-1' port_id can be
+ // Checks if a node by name is in the graph view.
+ bool HasNode(absl::string_view node_name) const {
+ return GetNode(node_name) != nullptr;
+ }
+
+ // Gets the specified input port. Note that the special '-1' port_id can be
// used to access the controlling nodes (i.e. the nodes connected to node_name
// through an incoming control dependency).
InputPort GetInputPort(absl::string_view node_name, int port_id) const {
return InputPort(GetNode(node_name), port_id);
}
- // Get the specified output port. Note that the special '-1' port_id can be
+ // Gets the specified output port. Note that the special '-1' port_id can be
// used to access the controlled nodes (i.e. the nodes connected to node_name
// through an outgoing control dependency).
OutputPort GetOutputPort(absl::string_view node_name, int port_id) const {
return OutputPort(GetNode(node_name), port_id);
}
- // Get the input (resp. output) port(s) in the immediate fanout (resp. fanin)
- // of an output (resp. input) port.
+ // Gets the input port(s) in the immediate fanout of an output port.
const absl::flat_hash_set<InputPort>& GetFanout(
const OutputPort& port) const {
return gtl::FindWithDefault(fanouts_, port, fanout_not_found_value_);
}
+ // Gets the output port(s) in the immediate fanin of an input port.
absl::flat_hash_set<OutputPort> GetFanin(const InputPort& port) const {
if (port.port_id >= 0) return {GetRegularFanin(port)};
@@ -162,9 +167,22 @@
return GetOutputPort(tensor_id.node(), tensor_id.index());
}
- // Get all the input (resp. output) ports in the immediate fanout (resp
- // fanin) of a node. Include the controlling nodes iff
- // include_controlling_nodes is true.
+ // Checks if a tensor id is a fanin of the node.
+ bool HasFanin(const NodeDef& node, const TensorId& fanin) const {
+ if (fanin.index() < -1) {
+ return false;
+ }
+ string fanin_string = TensorIdToString(fanin);
+ for (int i = 0; i < node.input_size(); ++i) {
+ if (node.input(i) == fanin_string) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ // Gets all the input ports in the immediate fanout of a node. Include the
+ // controlled nodes iff include_controlled_nodes is true.
absl::flat_hash_set<InputPort> GetFanouts(
const NodeDef& node, bool include_controlled_nodes) const {
absl::flat_hash_set<InputPort> result;
@@ -185,6 +203,8 @@
return result;
}
+ // Gets all the output ports in the immediate fanin of a node. Include the
+ // controlling nodes iff include_controlling_nodes is true.
absl::flat_hash_set<OutputPort> GetFanins(
const NodeDef& node, bool include_controlling_nodes) const {
absl::flat_hash_set<OutputPort> result;
@@ -198,7 +218,7 @@
return result;
}
- // Get the number of ports in the immediate fanin of a node. Count the
+ // Gets the number of ports in the immediate fanin of a node. Count the
// controlling nodes iff include_controlling_nodes is true.
int NumFanins(const NodeDef& node, bool include_controlling_nodes) const {
int count = 0;
@@ -211,14 +231,14 @@
return count;
}
- // Get the number of ports in the immediate fanout of a node. Count the
- // controlling nodes iff include_controlling_nodes is true.
- int NumFanouts(const NodeDef& node, bool include_controlling_nodes) const {
+ // Gets the number of ports in the immediate fanout of a node. Count the
+ // controlled nodes iff include_controlled_nodes is true.
+ int NumFanouts(const NodeDef& node, bool include_controlled_nodes) const {
int count = 0;
OutputPort port;
port.node = const_cast<NodeDefT*>(&node);
- const int first_port_id = include_controlling_nodes ? -1 : 0;
+ const int first_port_id = include_controlled_nodes ? -1 : 0;
const int last_port_id =
gtl::FindWithDefault(max_regular_output_port_, port.node, -1);
@@ -231,8 +251,8 @@
return count;
}
- // Get all the edges in the immediate fanout (resp fanin) of a node.
- // Include the control edges iff include_controlling_edges is true.
+ // Gets all the edges in the immediate fanout of a node. Include the
+ // controlled edges iff include_controlled_edges is true.
absl::flat_hash_set<Edge> GetFanoutEdges(
const NodeDef& node, bool include_controlled_edges) const {
absl::flat_hash_set<Edge> result;
@@ -248,14 +268,16 @@
auto it = fanouts_.find(port);
if (it != fanouts_.end()) {
for (auto itr = it->second.begin(); itr != it->second.end(); ++itr) {
- result.emplace(/*src*/ OutputPort(const_cast<NodeDefT*>(&node), i),
- /*dst*/ *itr);
+ result.emplace(/*src=*/OutputPort(const_cast<NodeDefT*>(&node), i),
+ /*dst=*/*itr);
}
}
}
return result;
}
+ // Gets all the edges in the immediate fanin of a node. Include the
+ // controlling edges iff include_controlling_edges is true.
absl::flat_hash_set<Edge> GetFaninEdges(
const NodeDef& node, bool include_controlling_edges) const {
absl::flat_hash_set<Edge> result;
@@ -265,8 +287,8 @@
auto it = nodes_.find(tensor_id.node());
if (it != nodes_.end()) {
- result.emplace(/*src*/ OutputPort(it->second, tensor_id.index()),
- /*dst*/ InputPort(const_cast<NodeDefT*>(&node), i));
+ result.emplace(/*src=*/OutputPort(it->second, tensor_id.index()),
+ /*dst=*/InputPort(const_cast<NodeDefT*>(&node), i));
}
}
return result;
diff --git a/tensorflow/core/grappler/graph_view_test.cc b/tensorflow/core/grappler/graph_view_test.cc
index cbf859a..404dcd3 100644
--- a/tensorflow/core/grappler/graph_view_test.cc
+++ b/tensorflow/core/grappler/graph_view_test.cc
@@ -230,6 +230,40 @@
EXPECT_EQ(0, (*fanin.begin()).port_id);
}
+TEST_F(GraphViewTest, HasNode) {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
+
+ GrapplerItem item;
+ TF_CHECK_OK(s.ToGraphDef(&item.graph));
+ GraphView graph(&item.graph);
+
+ EXPECT_EQ(true, graph.HasNode("a"));
+ EXPECT_EQ(false, graph.HasNode("b"));
+}
+
+TEST_F(GraphViewTest, HasFanin) {
+ tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+ Output a = ops::Const(s.WithOpName("a"), 0.0f, {10, 10});
+ Output b = ops::Square(s.WithOpName("b"), {a});
+ Output c = ops::Sqrt(s.WithOpName("c"), {b});
+ Output d = ops::AddN(s.WithOpName("d").WithControlDependencies(a), {b, c});
+
+ GrapplerItem item;
+ TF_CHECK_OK(s.ToGraphDef(&item.graph));
+ GraphView graph(&item.graph);
+
+ const NodeDef* d_node = graph.GetNode("d");
+ EXPECT_NE(nullptr, d_node);
+
+ EXPECT_EQ(true, graph.HasFanin(*d_node, {"a", Graph::kControlSlot}));
+ EXPECT_EQ(false, graph.HasFanin(*d_node, {"a", 0}));
+ EXPECT_EQ(true, graph.HasFanin(*d_node, {"b", 0}));
+ EXPECT_EQ(false, graph.HasFanin(*d_node, {"b", Graph::kControlSlot}));
+ EXPECT_EQ(true, graph.HasFanin(*d_node, {"c", 0}));
+ EXPECT_EQ(false, graph.HasFanin(*d_node, {"c", Graph::kControlSlot}));
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/grappler_item.cc b/tensorflow/core/grappler/grappler_item.cc
index f7cda35..2d71ac5 100644
--- a/tensorflow/core/grappler/grappler_item.cc
+++ b/tensorflow/core/grappler/grappler_item.cc
@@ -30,20 +30,22 @@
namespace tensorflow {
namespace grappler {
-GrapplerItem::GrapplerItem(const GrapplerItem& other, GraphDef* graph_def) {
- id = other.id;
- feed = other.feed;
- fetch = other.fetch;
- init_ops = other.init_ops;
- keep_ops = other.keep_ops;
- expected_init_time = other.expected_init_time;
- save_op = other.save_op;
- restore_op = other.restore_op;
- save_restore_loc_tensor = other.save_restore_loc_tensor;
- queue_runners = other.queue_runners;
- devices_ = other.devices_;
- allowed_optimizations_ = other.allowed_optimizations_;
- graph.Swap(graph_def);
+GrapplerItem GrapplerItem::WithGraph(GraphDef&& graph_def) const {
+ GrapplerItem item;
+ item.id = id;
+ item.feed = feed;
+ item.fetch = fetch;
+ item.init_ops = init_ops;
+ item.keep_ops = keep_ops;
+ item.expected_init_time = expected_init_time;
+ item.save_op = save_op;
+ item.restore_op = restore_op;
+ item.save_restore_loc_tensor = save_restore_loc_tensor;
+ item.queue_runners = queue_runners;
+ item.devices_ = devices_;
+ item.allowed_optimizations_ = allowed_optimizations_;
+ item.graph.Swap(&graph_def);
+ return item;
}
std::vector<const NodeDef*> GrapplerItem::MainOpsFanin() const {
@@ -112,6 +114,15 @@
result.insert(NodeName(queue_runner.cancel_op_name()));
}
}
+
+ if (!allowed_optimizations_.prune_ops_with_side_effects) {
+ for (const NodeDef& node : graph.node()) {
+ if (!IsFreeOfSideEffect(node)) {
+ result.insert(node.name());
+ }
+ }
+ }
+
return result;
}
diff --git a/tensorflow/core/grappler/grappler_item.h b/tensorflow/core/grappler/grappler_item.h
index 6ef4f14..1ae551f 100644
--- a/tensorflow/core/grappler/grappler_item.h
+++ b/tensorflow/core/grappler/grappler_item.h
@@ -35,12 +35,15 @@
// nodes, and potentially a set of nodes to feed.
struct GrapplerItem {
GrapplerItem() = default;
- GrapplerItem(const GrapplerItem& other, GraphDef&& graph_def)
- : GrapplerItem(other, &graph_def) {}
- // Swaps *graph_def with an empty GraphDef.
- GrapplerItem(const GrapplerItem& other, GraphDef* graph_def);
+ GrapplerItem(const GrapplerItem& other) = default;
+ GrapplerItem(GrapplerItem&& other) = default;
+ GrapplerItem& operator=(const GrapplerItem& other) = default;
+ GrapplerItem& operator=(GrapplerItem&& other) = default;
virtual ~GrapplerItem() = default;
+ // Create a copy of this GrapplerItem with graph swapped with the argument.
+ GrapplerItem WithGraph(GraphDef&& graph) const;
+
string id; // A unique id for this item
// Inputs
@@ -83,13 +86,12 @@
// Is it allowed to add nodes to the graph that do not have registered
// gradient function.
bool non_differentiable_rewrites = true;
- // By default we are not allowed to inline ops with side effects into the
- // main graph, because we can't guarantee that after pruning these ops will
- // be executed. However if we are optimizing a function library (see
- // meta_optimizer.cc) and a graph was instantiated by a function definition,
- // we can do that, because functions guarantee that all side effects will be
- // executed (see function_optimizer.cc for details).
- bool inline_ops_with_side_effects = false;
+
+ // By default we are allowed to prune ops with side-effects from the main
+ // graph if they are not in transitive fanin of the fetch nodes. If we are
+ // optimizing a graph that was instantiated by a function definition, we
+ // must keep all side effects intact.
+ bool prune_ops_with_side_effects = true;
};
const std::unordered_set<string>& devices() const;
diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 40eab8b..79578cb 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -385,7 +385,7 @@
srcs = [
"gpu_swapping_ops.cc",
],
- visibility = ["//tensorflow:__subpackages__"],
+ visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
@@ -461,7 +461,6 @@
"//tensorflow/core/grappler:devices",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:op_types",
- "//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/clusters:cluster",
"//tensorflow/core/grappler/costs:graph_properties",
"//tensorflow/core/grappler/costs:virtual_placer",
@@ -615,7 +614,6 @@
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:mutable_graph_view",
"//tensorflow/core/grappler:op_types",
- "//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/costs:graph_properties",
"//tensorflow/core/grappler/utils:frame",
"@com_google_absl//absl/container:flat_hash_set",
@@ -767,7 +765,6 @@
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:op_types",
- "//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/costs:graph_properties",
"//tensorflow/core/grappler/utils:frame",
],
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index e41b1cf..d35c00f 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -3561,8 +3561,7 @@
// Set up helper data structures.
nodes_to_preserve_ = item.NodesToPreserve();
fetch_nodes_known_ = !item.fetch.empty();
- *optimized_graph = item.graph;
- GrapplerItem optimized_item(item, optimized_graph);
+ GrapplerItem optimized_item(item);
optimized_graph_ = &optimized_item.graph;
node_map_.reset(new NodeMap(optimized_graph_));
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index 5e3e5d6..3882e3b 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -751,6 +751,12 @@
if (ModifiesFrameInfo(node)) {
return false;
}
+
+ // Removing LoopCond nodes can screw up the partitioner.
+ if (node.op() == "LoopCond") {
+ return false;
+ }
+
// Skip constants, they're already folded
if (IsConstant(node)) {
return false;
diff --git a/tensorflow/core/grappler/optimizers/data/map_and_batch_fusion.cc b/tensorflow/core/grappler/optimizers/data/map_and_batch_fusion.cc
index 9a7a751..84c4d82 100644
--- a/tensorflow/core/grappler/optimizers/data/map_and_batch_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_and_batch_fusion.cc
@@ -77,15 +77,22 @@
new_node.add_input(tmp->name());
}
- // Set `f` and `Targuments` attributes.
+ // Required attributes.
for (auto key : {"f", "Targuments"}) {
graph_utils::CopyAttribute(key, map_node, &new_node);
}
-
- // Set `output_types` and `output_shapes` attributes.
for (auto key : {"output_shapes", "output_types"}) {
graph_utils::CopyAttribute(key, batch_node, &new_node);
}
+
+ // Optional attributes.
+ // TODO(jsimsa): Support `use_inter_op_parallelism` and `sloppy`.
+ for (auto key : {"preserve_cardinality"}) {
+ if (gtl::FindOrNull(map_node.attr(), key)) {
+ graph_utils::CopyAttribute(key, map_node, &new_node);
+ }
+ }
+
return new_node;
}
diff --git a/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion.cc b/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion.cc
index 2b0a347..233d796 100644
--- a/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_and_filter_fusion.cc
@@ -55,8 +55,9 @@
}
// Optional attrs.
- for (auto key : {"use_inter_op_parallelism", "sloppy"}) {
- if (const auto* attr = gtl::FindOrNull(map_node.attr(), key)) {
+ for (auto key :
+ {"use_inter_op_parallelism", "sloppy", "preserve_cardinality"}) {
+ if (gtl::FindOrNull(map_node.attr(), key)) {
graph_utils::CopyAttribute(key, map_node, &fused_node);
}
}
diff --git a/tensorflow/core/grappler/optimizers/data/map_fusion.cc b/tensorflow/core/grappler/optimizers/data/map_fusion.cc
index 6ca0da2..6b8015f 100644
--- a/tensorflow/core/grappler/optimizers/data/map_fusion.cc
+++ b/tensorflow/core/grappler/optimizers/data/map_fusion.cc
@@ -62,9 +62,16 @@
gtl::FindOrNull(map_node.attr(), "use_inter_op_parallelism");
// Some graphs cannot execute with use_inter_op_parallelism=False, so we need
// to set it to true if one of the ops have it set to true.
- if (value_or_false(first_parallelism) || value_or_false(second_parallelism)) {
- (*fused_node.mutable_attr())["use_inter_op_parallelism"].set_b(true);
- }
+ (*fused_node.mutable_attr())["use_inter_op_parallelism"].set_b(
+ value_or_false(first_parallelism) || value_or_false(second_parallelism));
+
+ const auto* first_cardinality =
+ gtl::FindOrNull(parent_map_node.attr(), "preserve_cardinality");
+ const auto* second_cardinality =
+ gtl::FindOrNull(map_node.attr(), "preserve_cardinality");
+ (*fused_node.mutable_attr())["preserve_cardinality"].set_b(
+ value_or_false(first_cardinality) && value_or_false(second_cardinality));
+
return fused_node;
}
diff --git a/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc b/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc
index e1ac776..e330835 100644
--- a/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc
+++ b/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc
@@ -127,7 +127,7 @@
test::AsScalar<float>(4.0f));
TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
const auto twice_boosted_tensor = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(twice_boosted_tensor[0],
test::AsScalar<float>(2.0f));
@@ -223,7 +223,7 @@
test::AsScalar<float>(4.0f));
TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
const auto twice_boosted_tensor = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(twice_boosted_tensor[0],
test::AsScalar<float>(2.0f));
diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index 6968540..7069e5e 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -1307,6 +1307,9 @@
std::back_inserter(happens_before),
[](const GraphView::OutputPort port) { return port.node->name(); });
+ VLOG(3) << "Happens before set (size = " << happens_before.size()
+ << "): " << absl::StrJoin(happens_before, ", ");
+
// Nodes that must observe side effects to the captured resources.
std::vector<string> happens_after;
absl::c_transform(
@@ -1314,6 +1317,9 @@
std::back_inserter(happens_after),
[](const GraphView::InputPort port) { return port.node->name(); });
+ VLOG(3) << "Happens after set (size = " << happens_after.size()
+ << "): " << absl::StrJoin(happens_after, ", ");
+
// Regular (positional) inputs to the function call.
std::vector<SafeTensorId> inputs;
for (const string& input : func_node.input()) {
@@ -1466,7 +1472,7 @@
// for the function body, because functions have strict semantics.
if (num_fanouts == 0 && happens_after.empty() &&
- !ctx->allowed_optimizations().inline_ops_with_side_effects) {
+ ctx->allowed_optimizations().prune_ops_with_side_effects) {
return errors::Internal(
"Can't inline a function with a side-effectful op with empty "
"fanouts and empty output control edge set. Function body node: ",
@@ -1665,7 +1671,7 @@
gtl::FlatSet<string> add_ctrl_inputs;
// Remove all invalidated control inputs.
- for (int idx = 0; idx < node.input_size(); ++idx) {
+ for (int idx = 0; idx < node.input_size(); /* see below */) {
// TODO(ezhulenev): Use non-allocating TensorId after migrating
// `control_overrides()` to absl::flat_hash_set.
SafeTensorId input_tensor = ParseTensorName(node.input(idx));
@@ -1685,6 +1691,10 @@
for (const string& override : overrides->second) {
add_ctrl_inputs.insert(AsControlDependency(override));
}
+ } else {
+ // Go to the next input only if the current one was not invalidated,
+ // otherwise we need to check the swapped input as well.
+ ++idx;
}
}
diff --git a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
index 93a2fcd..c971eec 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
@@ -108,7 +108,7 @@
item.fetch = {"z"};
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -184,7 +184,7 @@
item.fetch = {"z1"};
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -284,7 +284,7 @@
item.fetch = {"z"};
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -368,7 +368,7 @@
item.fetch = {"z"};
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -418,7 +418,7 @@
item.feed.emplace_back("x4", test::AsScalar<float>(-1.0f));
item.feed.emplace_back("x3", test::AsScalar<int>(1234));
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
test::ExpectTensorEqual<float>(tensors_expected[1], tensors[1]);
@@ -549,7 +549,7 @@
item.feed.emplace_back("a", test::AsScalar<float>(2.0f));
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
@@ -748,7 +748,7 @@
item.feed.emplace_back("a", pi);
item.feed.emplace_back("b", pi);
- GrapplerItem optimized(item, std::move(optimized_graph));
+ GrapplerItem optimized = item.WithGraph(std::move(optimized_graph));
auto tensors_expected = EvaluateFetchNodes(item);
auto tensors = EvaluateFetchNodes(optimized);
ASSERT_EQ(tensors_expected.size(), 1);
@@ -812,8 +812,8 @@
// Return result of multiplication and a current value of the variable.
NDef("out_1", "Identity", {"f2"}, {{"T", DT_FLOAT}}, kDevice),
- NDef("out_2", "ReadVariableOp", {"v", "^f2"}, {{"dtype", DT_FLOAT}},
- kDevice)},
+ NDef("out_2", "ReadVariableOp", {"v", "^f1", "^f2"},
+ {{"dtype", DT_FLOAT}}, kDevice)},
// Function library.
{mul_func});
@@ -860,8 +860,8 @@
// Return values read directly from inlined nodes.
NDef("out_1", "Identity", {"f2/mul:0"}, {{"T", DT_FLOAT}}, kDevice),
- NDef("out_2", "ReadVariableOp", {"v", "^f2/add"}, {{"dtype", DT_FLOAT}},
- kDevice)},
+ NDef("out_2", "ReadVariableOp", {"v", "^f1/add", "^f2/add"},
+ {{"dtype", DT_FLOAT}}, kDevice)},
// Function library.
{mul_func});
@@ -876,9 +876,11 @@
EXPECT_EQ(tensors_expected[0].flat<float>()(0), 4.0); // mul
EXPECT_EQ(tensors_expected[1].flat<float>()(0), 3.0); // read variable
- GrapplerItem optimized(item, std::move(optimized_graph));
+ GrapplerItem optimized = item.WithGraph(std::move(optimized_graph));
auto tensors = EvaluateFetchNodes(optimized);
+ ASSERT_EQ(tensors.size(), 2);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+ test::ExpectTensorEqual<float>(tensors_expected[1], tensors[1]);
}
TEST_F(FunctionOptimizerTest, InlineIndirectFunctionWithDevicePlacement) {
@@ -934,6 +936,94 @@
CompareGraphs(expected, optimized_graph);
}
+TEST_F(FunctionOptimizerTest, InlineIndirectFunctionWithoutSideEffects) {
+ using test::function::NDef;
+ using FDH = FunctionDefHelper;
+
+ FunctionOptimizer optimizer(RewriterConfig::AGGRESSIVE);
+
+ const Tensor kOne = test::AsScalar<float>(1.0);
+ const Tensor kTwo = test::AsScalar<float>(2.0);
+ const TensorShape scalar = TensorShape({});
+
+ // MyMul doesn't have any side-effectful nodes in the function body, but the
+ // optimized graph has a control dependency edge `f1->f2`.
+ FunctionDef mul_func = FunctionDefHelper::Create(
+ "MyMul", {"x:T", "y:T"}, {"z:T"}, {"T: {float, double}"},
+ {{{"mul"}, "Mul", {"x", "y"}, {{"T", "$T"}}}},
+ /* Mapping between function returns and function node outputs. */
+ {{"z", "mul:z:0"}});
+
+ // Build a graph to compute:
+ // a = Placeholder
+ // b = Placeholder
+ // f1 = MyMul(a, b)
+ // f2 = MyMul(a, b, ^f1) <-- control dependency on inlined function!
+ // return f2
+ GrapplerItem item;
+ item.fetch = {"out"};
+ item.graph = test::function::GDef(
+ {NDef("a", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+ NDef("b", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+
+ // Call function first time.
+ NDef("f1", "PartitionedCall", {"a", "b"},
+ {{"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+ {"Tout", DataTypeSlice{DT_FLOAT}},
+ {"f", FDH::FunctionRef("MyMul", {{"T", DT_FLOAT}})}},
+ kDevice),
+
+ // Call function second time.
+ NDef("f2", "PartitionedCall", {"f1", "f1", "^f1"},
+ {{"Tin", DataTypeSlice{DT_FLOAT, DT_FLOAT}},
+ {"Tout", DataTypeSlice{DT_FLOAT}},
+ {"f", FDH::FunctionRef("MyMul", {{"T", DT_FLOAT}})}},
+ kDevice),
+
+ // Return result of f2.
+ NDef("out", "Identity", {"f2"}, {{"T", DT_FLOAT}}, kDevice)},
+
+ // Function library.
+ {mul_func});
+
+ GraphDef optimized_graph;
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &optimized_graph));
+
+ GraphDef expected = test::function::GDef(
+ {NDef("a", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+ NDef("b", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+
+ // Function body of a first function call inlined into the graph.
+ NDef("f1/x", "Identity", {"a:0"}, {{"T", DT_FLOAT}}, kDevice),
+ NDef("f1/y", "Identity", {"b:0"}, {{"T", DT_FLOAT}}, kDevice),
+ NDef("f1/mul", "Mul", {"f1/x", "f1/y"}, {{"T", DT_FLOAT}}, kDevice),
+
+ // Function body of a second function call also inlined into the graph,
+ // and input nodes read directly from the inlined nodes of the first
+ // function call, and control dependency edge removed.
+ NDef("f2/x", "Identity", {"f1/mul:0"}, {{"T", DT_FLOAT}}, kDevice),
+ NDef("f2/y", "Identity", {"f1/mul:0"}, {{"T", DT_FLOAT}}, kDevice),
+ NDef("f2/mul", "Mul", {"f2/x", "f2/y"}, {{"T", DT_FLOAT}}, kDevice),
+
+ // Return directly from inlined node of f2.
+ NDef("out", "Identity", {"f2/mul:0"}, {{"T", DT_FLOAT}}, kDevice)},
+
+ // Function library.
+ {mul_func});
+
+ CompareGraphs(expected, optimized_graph);
+
+ item.feed.emplace_back("a", kOne);
+ item.feed.emplace_back("b", kTwo);
+
+ auto tensors_expected = EvaluateFetchNodes(item);
+ ASSERT_EQ(tensors_expected.size(), 1);
+
+ GrapplerItem optimized = item.WithGraph(std::move(optimized_graph));
+ auto tensors = EvaluateFetchNodes(optimized);
+ test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+}
+
TEST_F(FunctionOptimizerTest, SpecializeFunctionXTimesTwo) {
using test::function::NDef;
@@ -977,7 +1067,7 @@
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -1041,7 +1131,7 @@
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -1104,7 +1194,7 @@
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -1184,7 +1274,7 @@
item.feed.emplace_back("x", pi);
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -1300,7 +1390,7 @@
item.feed = {{"xf", pi}, {"yf", pi}, {"xi", four}, {"yi", four}};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
@@ -1409,7 +1499,7 @@
item.feed = {{"xf", pi}, {"yf", pi}};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
ASSERT_EQ(tensors_expected.size(), tensors.size());
@@ -1570,7 +1660,7 @@
item.feed = {{"xf", pi}, {"yf", pi}};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
ASSERT_EQ(tensors_expected.size(), tensors.size());
diff --git a/tensorflow/core/grappler/optimizers/layout_optimizer.cc b/tensorflow/core/grappler/optimizers/layout_optimizer.cc
index f465350..8f25a1c 100644
--- a/tensorflow/core/grappler/optimizers/layout_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/layout_optimizer.cc
@@ -29,7 +29,6 @@
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/optimizers/layout_optimizer.h"
-#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/frame.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/str_util.h"
@@ -1969,9 +1968,9 @@
// Expand all nodes which is in NHWC, but supports NCHW or is layout agnostic.
Status Expand() {
int node_size_original = graph_->node_size();
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
- TF_RETURN_IF_ERROR(IdentifyFrames(*graph_, &frames, &num_frames));
+
+ FrameView frame_view;
+ TF_RETURN_IF_ERROR(frame_view.InferFromGraph(*graph_));
// This is the first pass where we expand the nodes which support NCHW.
std::set<string> ops_format_supported = GetOpsFormatSupported();
@@ -1983,7 +1982,7 @@
if (ops_format_supported.find(graph_->node(i).op()) !=
ops_format_supported.end()) {
auto node = graph_->mutable_node(i);
- bool is_in_frame = !frames[node].empty();
+ bool is_in_frame = frame_view.IsInFrame(*node);
OptimizeContext opt_cxt(graph_, node, node_map_, graph_properties_,
virtual_placer_, nodes_to_preserve_,
is_in_frame);
@@ -2033,7 +2032,7 @@
if (ops_format_agnostic.find(graph_->node(i).op()) !=
ops_format_agnostic.end()) {
auto node = graph_->mutable_node(i);
- bool is_in_frame = !frames[node].empty();
+ bool is_in_frame = frame_view.IsInFrame(*node);
OptimizeContext opt_cxt(graph_, node, node_map_, graph_properties_,
virtual_placer_, nodes_to_preserve_,
is_in_frame);
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer.cc b/tensorflow/core/grappler/optimizers/loop_optimizer.cc
index 775fb9a..3606473 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer.cc
@@ -35,7 +35,6 @@
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/optimizers/constant_folding.h"
#include "tensorflow/core/grappler/optimizers/evaluation_utils.h"
-#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/frame.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
@@ -380,14 +379,14 @@
Status LoopInvariantNodeMotionOptimizer::Optimize() {
node_map_.reset(new NodeMap(optimized_graph_));
- FrameMap frame_map;
- int num_frames;
- TF_RETURN_IF_ERROR(IdentifyFramesWithNodeMap(*optimized_graph_, *node_map_,
- &frame_map, &num_frames));
+ FrameView frame_view;
+ // TODO(ezhulenev): Use GraphView when migrated from NodeMap.
+ TF_RETURN_IF_ERROR(frame_view.InferFromGraph(*optimized_graph_));
+
std::deque<int> worklist;
- for (auto iter = frame_map.begin(); iter != frame_map.end(); ++iter) {
- auto* node = iter->first;
- auto& frame_ids = iter->second;
+ for (const NodeDef& node : optimized_graph_->node()) {
+ const std::vector<int>& frame_ids = frame_view.Frames(node);
+
if (frame_ids.size() >= 3) {
for (unsigned int i = 1; i < frame_ids.size() - 1; ++i) {
frame_parent_[frame_ids[i]] = frame_ids[i - 1];
@@ -400,18 +399,18 @@
}
if (!frame_ids.empty()) {
frame_children_.insert(std::make_pair(frame_ids.back(), empty_set_));
- if (node->op() == "LoopCond") {
+ if (node.op() == "LoopCond") {
if (loop_cond_.count(frame_ids.back())) {
return errors::InvalidArgument(
"Loop ", frame_ids.back(),
- " has more than one LoopCond node: ", node->name(), " and ",
+ " has more than one LoopCond node: ", node.name(), " and ",
loop_cond_[frame_ids.back()]->name());
}
- loop_cond_[frame_ids.back()] = node;
+ loop_cond_[frame_ids.back()] = &node;
}
- if (IsEnter(*node) && node->attr().at("is_constant").b()) {
+ if (IsEnter(node) && node.attr().at("is_constant").b()) {
invariant_enters_[frame_ids.back()].push_back(
- const_cast<NodeDef*>(node));
+ const_cast<NodeDef*>(&node));
}
}
}
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer.h b/tensorflow/core/grappler/optimizers/loop_optimizer.h
index 7c04f55..d467237 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer.h
@@ -19,7 +19,6 @@
#include <unordered_set>
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/optimizers/graph_optimizer.h"
-#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/frame.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc b/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc
index 7332839..587767c 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc
@@ -101,27 +101,30 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).back(), 0);
+ EXPECT_EQ(frames.num_frames(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).back(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("VariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("VariantAdd")).back(), 0);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).back(), 0);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("VariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("VariantAdd")).back(), 0);
+ }
}
TEST_F(LoopOptimizerTest, Const) {
@@ -149,26 +152,29 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("Const")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("Const")).back(), 0);
+ EXPECT_EQ(frames.num_frames(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).back(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("Const")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("Const")).back(), 0);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("Const")).size(), 0);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("Const")).size(), 0);
+ }
}
TEST_F(LoopOptimizerTest, ControlOutput) {
@@ -197,24 +203,27 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+ EXPECT_EQ(frames.num_frames(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).back(), 0);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).back(), 0);
+ }
}
TEST_F(LoopOptimizerTest, NestedLoop1) {
@@ -258,31 +267,34 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).back(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).back(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).back(), 0);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 0);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).back(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).back(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd")).size(), 0);
+ }
}
TEST_F(LoopOptimizerTest, NestedLoop2) {
@@ -326,27 +338,30 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).back(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).back(), 1);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("VariantAdd2")).back(), 1);
+ }
}
TEST_F(LoopOptimizerTest, NestedLoopConst1) {
@@ -390,28 +405,31 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).back(), 1);
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).back(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("Const2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("Const2")).back(), 1);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).back(), 0);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).back(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("Const2")).size(), 1);
+ EXPECT_EQ(frames.Frames(*view.GetNode("Const2")).back(), 0);
+ }
}
TEST_F(LoopOptimizerTest, NestedLoopConst2) {
@@ -455,26 +473,29 @@
LoopOptimizer optimizer;
EnableOnlyLoopInvariantNodeMotion(&optimizer);
GraphDef output;
- Status status = optimizer.Optimize(nullptr, item, &output);
- TF_EXPECT_OK(status);
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
- std::unique_ptr<NodeMap> node_map;
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
+ { // Original graph.
+ GraphView view(&graph);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
- node_map.reset(new NodeMap(&graph));
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 2);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).back(), 1);
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).back(), 1);
+ ASSERT_EQ(frames.Frames(*view.GetNode("Const2")).size(), 2);
+ EXPECT_EQ(frames.Frames(*view.GetNode("Const2")).back(), 1);
+ }
- node_map.reset(new NodeMap(&output));
- EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 0);
- EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 0);
+ { // Optimized graph.
+ GraphView view(&output);
+ FrameView frames;
+ TF_EXPECT_OK(frames.InferFromGraphView(view));
+
+ EXPECT_EQ(frames.num_frames(), 2);
+ ASSERT_EQ(frames.Frames(*view.GetNode("InvariantAdd2")).size(), 0);
+ ASSERT_EQ(frames.Frames(*view.GetNode("Const2")).size(), 0);
+ }
}
void VerifyGraphsEqual(const GraphDef& original_graph,
diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer.cc b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
index 453db5d..227c2bb 100644
--- a/tensorflow/core/grappler/optimizers/memory_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
@@ -1306,13 +1306,12 @@
Status MemoryOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
GraphDef* optimized_graph) {
- *optimized_graph = item.graph;
+ GrapplerItem optimized_item(item);
RecomputationRewritingPass(optimization_level_,
- recomputation_targets_name_scope_, optimized_graph,
- item);
+ recomputation_targets_name_scope_,
+ &optimized_item.graph, item);
- GrapplerItem optimized_item(item, optimized_graph);
std::unordered_set<string> skip_list;
// Bound the number of rewrite passes to avoid long processing times on graphs
// that simply won't fit in memory.
diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc b/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc
index 75285b0..356b23d 100644
--- a/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc
@@ -279,7 +279,7 @@
EXPECT_EQ("^swap_out_e_0", new_c.input(1));
// Run the optimizer a second time to ensure it's idempotent.
- GrapplerItem item_copy(item, std::move(output));
+ GrapplerItem item_copy = item.WithGraph(std::move(output));
status = optimizer.Optimize(cluster.get(), item_copy, &output);
TF_EXPECT_OK(status);
@@ -287,7 +287,7 @@
item.fetch = {"e"};
item.init_ops = {init.name()};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
#endif
@@ -337,7 +337,7 @@
#if GOOGLE_CUDA
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
for (int i = 0; i < item.fetch.size(); ++i) {
test::ExpectTensorEqual<float>(tensors_expected[i], tensors[i]);
@@ -386,7 +386,7 @@
#if GOOGLE_CUDA
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
#endif
@@ -474,7 +474,7 @@
item.fetch = {"exp"};
item.init_ops = {"variable"};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
@@ -505,7 +505,7 @@
item.fetch = {"exp"};
item.init_ops = {"variable"};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
#endif
@@ -598,7 +598,7 @@
item.fetch = {"assign0", "assign1"};
item.init_ops = {"exp1", "variable1"};
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
for (int i = 0; i < tensors_expected.size(); ++i) {
test::ExpectTensorEqual<float>(tensors_expected[i], tensors[i]);
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index 7c83036..7b788c6 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -440,7 +440,7 @@
item.graph)
.ToProto();
- GrapplerItem trimmed_item(item, std::move(trimmed_graph));
+ GrapplerItem trimmed_item = item.WithGraph(std::move(trimmed_graph));
VLOG(1) << absl::Substitute(
"Deleted $0 unreachable functions from the graph (library size = $1)",
@@ -533,9 +533,10 @@
VLOG(3) << added_devices.error_message();
}
- // We can safely inline nested function calls with side-effectful ops into
- // the function body (see function_optimizer.cc for details).
- func_item.allowed_optimizations().inline_ops_with_side_effects = true;
+ // We are not allowed to prune side effects from the graph instantiated
+ // by the function definition, because we must guarantee function
+ // execution semantics wrt side effects (see function_optimizer.cc).
+ func_item.allowed_optimizations().prune_ops_with_side_effects = false;
// Optimize function body graph.
GraphDef optimized_func_graph;
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
index 42b867b..12db5d6 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
@@ -396,7 +396,7 @@
item.feed.emplace_back("b", test::AsScalar<int>(4));
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
@@ -502,7 +502,7 @@
item.feed.emplace_back("b", test::AsScalar<float>(3.123f));
auto tensors_expected = EvaluateFetchNodes(item);
- GrapplerItem optimized(item, std::move(output));
+ GrapplerItem optimized = item.WithGraph(std::move(output));
auto tensors = EvaluateFetchNodes(optimized);
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc
index d8e62e0..3fb3f2b 100644
--- a/tensorflow/core/grappler/optimizers/remapper.cc
+++ b/tensorflow/core/grappler/optimizers/remapper.cc
@@ -665,7 +665,7 @@
std::reverse(topo_sorted_graph.mutable_node()->begin(),
topo_sorted_graph.mutable_node()->end());
- GrapplerItem topo_sorted_item(item, std::move(topo_sorted_graph));
+ GrapplerItem topo_sorted_item = item.WithGraph(std::move(topo_sorted_graph));
RemapperContext ctx(topo_sorted_item);
// Skip nodes that were invalidated by a remapper, e.g. do not process BiasAdd
diff --git a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
index 0d4aaf6..e537b3d 100644
--- a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
@@ -790,20 +790,17 @@
return root;
}
-void PartitionByLoopStructure(const FrameMap& frame_map,
+void PartitionByLoopStructure(const FrameView& frame_view,
std::vector<NodeDef*> nodes,
std::vector<std::vector<NodeDef*>>* loop_groups) {
// It is assumed that two nodes with identical loop containment have
- // identical integer vectors. Represent those by 64 bit hashes.
+ // identical integer vectors. Represent those by 64 bit hashes.
std::unordered_map<uint64, std::vector<NodeDef*>> loop_sets;
for (NodeDef* nd : nodes) {
uint64 hash = 0;
- const auto& it = frame_map.find(nd);
- if (it != frame_map.end()) {
- const std::vector<int>& loop_ids = it->second;
- for (int id : loop_ids) {
- hash = Hash64Combine(hash, static_cast<uint64>(id));
- }
+ const std::vector<int>& loop_ids = frame_view.Frames(*nd);
+ for (int id : loop_ids) {
+ hash = Hash64Combine(hash, static_cast<uint64>(id));
}
loop_sets[hash].push_back(nd);
}
@@ -821,10 +818,11 @@
GraphOpOccurrences occ;
FindOpOccurrences(graph, op_name_set_, &occ);
if (!occ.empty()) {
- FrameMap frame_map;
- int num_frames;
- LOG_WARNING_AND_RETURN_IF_ERROR(
- IdentifyFramesWithNodeMap(*graph, *node_map_, &frame_map, &num_frames));
+ FrameView frame_view;
+ // TODO(ezhulenev): Pass a GraphView when this optimizer will be migrated
+ // from NodeMap.
+ LOG_WARNING_AND_RETURN_IF_ERROR(frame_view.InferFromGraph(*graph));
+
for (auto& dt : occ) {
VLOG(2) << "Processing device " << dt.first;
const DevOpOccurrences& dev_occ = dt.second;
@@ -841,26 +839,26 @@
// Nodes with a common depth and root path are now grouped
// in the same Tree struct. Split those groups into subgroups that
// share identical loop nesting.
- status = ApplyToAll(
- root.get(), [this, rewriter, graph, &frame_map, &op_name](Tree* t) {
- VLOG(2) << "applied to tree node " << t->edge_ << " at depth "
- << t->depth_ << " of size " << t->nodes_.size();
- if (t->nodes_.size() > 1) {
- std::vector<std::vector<NodeDef*>> loop_groups;
- PartitionByLoopStructure(frame_map, t->nodes_, &loop_groups);
- for (auto& lg : loop_groups) {
- if (lg.size() > 1) {
- bool applied = false;
- Status s = OrderNodeSet(&lg);
- TF_RETURN_IF_ERROR(s);
- VLOG(1) << "Applying Rewriter for " << op_name;
- s = rewriter->Rewrite(this, graph, op_name, lg, &applied);
- LOG_WARNING_AND_RETURN_IF_ERROR(s);
- }
- }
+ status = ApplyToAll(root.get(), [this, rewriter, graph, &frame_view,
+ &op_name](Tree* t) {
+ VLOG(2) << "applied to tree node " << t->edge_ << " at depth "
+ << t->depth_ << " of size " << t->nodes_.size();
+ if (t->nodes_.size() > 1) {
+ std::vector<std::vector<NodeDef*>> loop_groups;
+ PartitionByLoopStructure(frame_view, t->nodes_, &loop_groups);
+ for (auto& lg : loop_groups) {
+ if (lg.size() > 1) {
+ bool applied = false;
+ Status s = OrderNodeSet(&lg);
+ TF_RETURN_IF_ERROR(s);
+ VLOG(1) << "Applying Rewriter for " << op_name;
+ s = rewriter->Rewrite(this, graph, op_name, lg, &applied);
+ LOG_WARNING_AND_RETURN_IF_ERROR(s);
}
- return Status::OK();
- });
+ }
+ }
+ return Status::OK();
+ });
if (!status.ok()) {
break;
}
diff --git a/tensorflow/core/grappler/utils.cc b/tensorflow/core/grappler/utils.cc
index 2977544..90ad04c 100644
--- a/tensorflow/core/grappler/utils.cc
+++ b/tensorflow/core/grappler/utils.cc
@@ -144,11 +144,16 @@
outputs.insert(nodes_[NodeName(new_output_name)]);
}
+string TensorIdToString(const TensorId& tensor_id) {
+ return tensor_id.index() == 0 ? string(tensor_id.node())
+ : tensor_id.ToString();
+}
+
bool IsSameInput(const string& name1, const string& name2) {
if (name1 == name2) return true;
TensorId tensor1 = ParseTensorName(name1);
TensorId tensor2 = ParseTensorName(name2);
- return tensor1.node() == tensor2.node() && tensor1.index() == tensor2.index();
+ return tensor1 == tensor2;
}
bool IsControlInput(const string& name) {
diff --git a/tensorflow/core/grappler/utils.h b/tensorflow/core/grappler/utils.h
index b1e2d4e..89a87af 100644
--- a/tensorflow/core/grappler/utils.h
+++ b/tensorflow/core/grappler/utils.h
@@ -100,6 +100,10 @@
std::vector<T> vector_;
};
+// Returns formatted string from TensorId specific to grappler. Specifically,
+// for the 0 port (first output), only the node name is returned.
+string TensorIdToString(const TensorId& tensor_id);
+
// True iff 'name' refers to a control inputs, i.e. a node name prefixed with
// the ^ character.
bool IsControlInput(const string& name);
diff --git a/tensorflow/core/grappler/utils/BUILD b/tensorflow/core/grappler/utils/BUILD
index 2b9448e..c0f19d3 100644
--- a/tensorflow/core/grappler/utils/BUILD
+++ b/tensorflow/core/grappler/utils/BUILD
@@ -74,8 +74,9 @@
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
+ "//tensorflow/core/grappler:graph_view",
"//tensorflow/core/grappler:op_types",
- "//tensorflow/core/grappler:utils",
+ "@com_google_absl//absl/container:flat_hash_map",
],
)
diff --git a/tensorflow/core/grappler/utils/frame.cc b/tensorflow/core/grappler/utils/frame.cc
index df5f4ff..2484b35 100644
--- a/tensorflow/core/grappler/utils/frame.cc
+++ b/tensorflow/core/grappler/utils/frame.cc
@@ -15,77 +15,128 @@
#include "tensorflow/core/grappler/utils/frame.h"
#include <deque>
-#include <stack>
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/grappler/graph_view.h"
#include "tensorflow/core/grappler/op_types.h"
-#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
namespace grappler {
-Status IdentifyFrames(const GraphDef& graph, FrameMap* frame_map,
- int* num_frames) {
- NodeMap node_map(const_cast<GraphDef*>(&graph));
- return IdentifyFramesWithNodeMap(graph, node_map, frame_map, num_frames);
-}
+namespace {} // namespace
-Status IdentifyFramesWithNodeMap(const GraphDef& graph, const NodeMap& node_map,
- FrameMap* frame_map, int* num_frames) {
- std::deque<std::pair<const NodeDef*, std::vector<int>>> ready_nodes;
- for (const NodeDef& node : graph.node()) {
+Status FrameView::InferFromGraphView(const GraphView& graph_view) {
+ if (is_inferred_) {
+ return errors::Internal("FrameView was already inferred from the graph");
+ }
+ is_inferred_ = true;
+
+ std::deque<const NodeDef*> ready_nodes;
+
+ // All nodes without inputs are automatically added to the ready queue.
+ for (const NodeDef& node : graph_view.graph()->node()) {
if (node.input_size() == 0) {
- std::vector<int> empty;
- ready_nodes.emplace_back(&node, empty);
- (*frame_map)[&node] = empty;
+ ready_nodes.push_back(&node);
+ node_to_frames_[&node] = node_has_no_frames_;
}
}
- std::map<string, int> name_to_id;
+
+ // We assign unique int id to each frame, and use this map to track what
+ // frames we've already seen in the graph.
+ absl::flat_hash_map<string, int> frame_name_to_id;
+
while (!ready_nodes.empty()) {
- auto ready_node = ready_nodes.front();
- for (const auto& fanout : node_map.GetOutputs(ready_node.first->name())) {
- if (frame_map->count(fanout) < 1) {
- std::vector<int> frame_ids = ready_node.second;
- if (IsExit(*ready_node.first)) {
+ const NodeDef* ready_node = ready_nodes.front();
+
+ absl::flat_hash_set<GraphView::InputPort> fanouts =
+ graph_view.GetFanouts(*ready_node, /*include_controlled_nodes=*/true);
+
+ for (const GraphView::InputPort& fanout : fanouts) {
+ if (node_to_frames_.count(fanout.node) < 1) {
+ // If we have never seen this node before, we add all frames from the
+ // incoming node (and pop/push frames if coming from Exit/Enter nodes).
+ std::vector<int> frame_ids = node_to_frames_[ready_node];
+
+ if (IsExit(*ready_node)) {
frame_ids.pop_back();
}
- if (IsEnter(*fanout)) {
- CHECK(fanout->attr().count("frame_name"))
- << "Missing frame name for the Enter node " << fanout->name();
- string name = fanout->attr().at("frame_name").s();
- int id;
- if (name_to_id.count(name)) {
- id = name_to_id[name];
- } else {
- id = name_to_id.size();
- name_to_id[name] = id;
+
+ if (IsEnter(*fanout.node)) {
+ const AttrValue* frame_name_attr =
+ AttrSlice(*fanout.node).Find("frame_name");
+
+ if (!frame_name_attr) {
+ return errors::InvalidArgument(
+ "Missing frame name for the Enter node: ",
+ SummarizeNodeDef(*fanout.node));
}
- frame_ids.push_back(id);
+
+ absl::string_view frame_name = frame_name_attr->s();
+ int frame_id;
+
+ if (frame_name_to_id.count(frame_name)) {
+ frame_id = frame_name_to_id[frame_name];
+ } else {
+ frame_id = static_cast<int>(frame_name_to_id.size());
+ frame_name_to_id[frame_name] = frame_id;
+ }
+
+ frame_ids.push_back(frame_id);
}
- ready_nodes.emplace_back(fanout, frame_ids);
- (*frame_map)[fanout] = frame_ids;
+
+ ready_nodes.push_back(fanout.node);
+ node_to_frames_[fanout.node] = std::move(frame_ids);
+
} else {
- auto frame_ids_fanout = (*frame_map)[fanout];
- auto frame_ids_node = ready_node.second;
- if (IsEnter(*fanout)) {
+ // If we've already seen this node before, we need to make sure that
+ // graph is correct and same nodes doesn't have incoming edges with
+ // conflicting frames (all inputs must be produces in the same frame).
+
+ std::vector<int> frame_ids_fanout = node_to_frames_[fanout.node];
+ std::vector<int> frame_ids_node = node_to_frames_[ready_node];
+
+ if (IsEnter(*fanout.node)) {
frame_ids_fanout.pop_back();
}
- if (IsExit(*ready_node.first)) {
+ if (IsExit(*ready_node)) {
frame_ids_node.pop_back();
}
+
if (frame_ids_node != frame_ids_fanout) {
return errors::InvalidArgument(
- "Invalid graph: Frame ids for node ", ready_node.first->name(),
- " does not match frame ids for it's fanout.");
+ "Invalid graph: Frame ids for node ", ready_node->name(),
+ " does not match frame ids for it's fanout ",
+ fanout.node->name());
}
}
}
+
ready_nodes.pop_front();
}
- *num_frames = name_to_id.size();
+
+ num_frames_ = static_cast<int>(frame_name_to_id.size());
return Status::OK();
}
+Status FrameView::InferFromGraph(const GraphDef& graph) {
+ return InferFromGraphView(GraphView(&graph));
+}
+
+const std::vector<int>& FrameView::Frames(const NodeDef& node) const {
+ DCHECK(is_inferred_) << "FrameView is not initialized";
+ auto frames = node_to_frames_.find(&node);
+ if (frames == node_to_frames_.end()) {
+ LOG(WARNING) << "Node doesn't belong to the graph used for initialization";
+ return node_has_no_frames_;
+ } else {
+ return frames->second;
+ }
+}
+
+bool FrameView::IsInFrame(const NodeDef& node) const {
+ return !Frames(node).empty();
+}
+
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/utils/frame.h b/tensorflow/core/grappler/utils/frame.h
index 95b7274..04c6588 100644
--- a/tensorflow/core/grappler/utils/frame.h
+++ b/tensorflow/core/grappler/utils/frame.h
@@ -17,25 +17,52 @@
#define TENSORFLOW_CORE_GRAPPLER_UTILS_FRAME_H_
#include <unordered_map>
+#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/framework/graph.pb.h"
-#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/grappler/graph_view.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace grappler {
-using FrameMap = std::unordered_map<const NodeDef*, std::vector<int>>;
+// FrameView is a helper class that allows to find in what execution frames (if
+// any) the given node can be running in. It's constructed from an immutable
+// GraphView, and any modification of the underlying graph might invalidate it.
+//
+// All execution frames assigned an unique integer id, but they do not have any
+// meaning whatsoever, it's just a sequence number.
+//
+// See the paper "Dynamic Control Flow in Large-Scale Machine Learning" for
+// detailed explanation of execution frames (https://arxiv.org/abs/1805.01772).
+class FrameView {
+ public:
+ FrameView() : is_inferred_(false), num_frames_(0) {}
-// Returns the number of frames present in the graph, and populates
-// the 'frames' argument with the collection of frames (denoted by their
-// frame ids) in the outermost-to-innermost order. Frame ids are arbitrary.
-Status IdentifyFrames(const GraphDef& graph, FrameMap* frame_map,
- int* num_frames);
+ // Infers nodes execution frames from the GraphView. Returns an error if
+ // called multiple times.
+ Status InferFromGraphView(const GraphView& graph_view);
+ // Infers nodes execution by constructing temporary GraphView and passing it
+ // to InferFromGraphView.
+ Status InferFromGraph(const GraphDef& graph);
-// As above, but use an existing NodeMap for graph instead of building it
-// from scratch.
-Status IdentifyFramesWithNodeMap(const GraphDef& graph, const NodeMap& node_map,
- FrameMap* frame_map, int* num_frames);
+ // Returns all frames of the given node (denoted by their frame ids) in
+ // outermost-to-innermost order.
+ const std::vector<int>& Frames(const NodeDef& node) const;
+
+ // Returns true iff the node is at least in one execution frame.
+ bool IsInFrame(const NodeDef& node) const;
+
+ int num_frames() const { return num_frames_; }
+ bool is_inferred() const { return is_inferred_; }
+
+ private:
+ bool is_inferred_; // true if it was inferred from the graph
+ int num_frames_; // number of frames present in a graph
+ absl::flat_hash_map<const NodeDef*, std::vector<int>> node_to_frames_;
+
+ // We return a reference to this vector if node has no frames.
+ const std::vector<int> node_has_no_frames_;
+};
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/utils/frame_test.cc b/tensorflow/core/grappler/utils/frame_test.cc
index df76083..cc82e0e 100644
--- a/tensorflow/core/grappler/utils/frame_test.cc
+++ b/tensorflow/core/grappler/utils/frame_test.cc
@@ -23,7 +23,7 @@
namespace grappler {
namespace {
-class IdentifyFramesTest : public ::testing::Test {
+class FrameViewTest : public ::testing::Test {
protected:
static NodeDef CreateNode(const string& name,
const std::vector<string>& inputs) {
@@ -53,19 +53,17 @@
}
};
-TEST_F(IdentifyFramesTest, NestedLoop) {
+TEST_F(FrameViewTest, NestedLoop) {
GraphDef graph;
// Create a two-level nested loop
*graph.add_node() = CreateNode("0", {});
- *graph.add_node() =
- CreateNode("1", "Enter", "map/while/while_context1", {"0"});
+ *graph.add_node() = CreateNode("1", "Enter", "while/context1", {"0"});
*graph.add_node() = CreateNode("2", {"1"});
*graph.add_node() = CreateNode("3", "Merge", {"2", "14"});
*graph.add_node() = CreateNode("4", {"3"});
*graph.add_node() = CreateNode("5", "Switch", {"4"});
*graph.add_node() = CreateNode("6", {"5"});
- *graph.add_node() =
- CreateNode("7", "Enter", "map/while/while_context2", {"6"});
+ *graph.add_node() = CreateNode("7", "Enter", "while/context2", {"6"});
*graph.add_node() = CreateNode("8", {"7"});
*graph.add_node() = CreateNode("9", "Merge", {"8", "12"});
*graph.add_node() = CreateNode("10", {"9"});
@@ -77,118 +75,106 @@
*graph.add_node() = CreateNode("16", "Exit", {"15"});
*graph.add_node() = CreateNode("17", {"16"});
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ FrameView frame_view;
+ ASSERT_TRUE(frame_view.InferFromGraph(graph).ok());
+
std::unordered_map<string, std::vector<int>> expected = {
{"0", {}}, {"1", {0}}, {"2", {0}}, {"3", {0}},
{"4", {0}}, {"5", {0}}, {"6", {0}}, {"7", {0, 1}},
{"8", {0, 1}}, {"9", {0, 1}}, {"10", {0, 1}}, {"11", {0, 1}},
{"12", {0, 1}}, {"13", {0, 1}}, {"14", {0}}, {"15", {0}},
{"16", {0}}, {"17", {}}};
- EXPECT_EQ(num_frames, 2);
- EXPECT_EQ(frames.size(), expected.size());
- std::cout << "Number of frame: " << num_frames << std::endl;
- for (const auto& node : frames) {
- std::cout << node.first->name() << ": ";
- EXPECT_EQ(node.second.size(), expected[node.first->name()].size());
- for (int i = 0; i < node.second.size(); i++) {
- EXPECT_EQ(expected[node.first->name()][i], node.second[i]);
- std::cout << node.second[i] << " ";
- }
- std::cout << std::endl;
+
+ EXPECT_EQ(frame_view.num_frames(), 2);
+ for (const NodeDef& node : graph.node()) {
+ std::vector<int> expected_frames = expected[node.name()];
+ std::vector<int> node_frames = frame_view.Frames(node);
+ EXPECT_EQ(expected_frames, node_frames);
}
}
-TEST_F(IdentifyFramesTest, MultipleInputsToEnter) {
+TEST_F(FrameViewTest, MultipleInputsToEnter) {
GraphDef graph;
*graph.add_node() = CreateNode("0", {});
*graph.add_node() = CreateNode("1", {});
- *graph.add_node() =
- CreateNode("2", "Enter", "map/while/while_context", {"0", "1"});
+ *graph.add_node() = CreateNode("2", "Enter", "while/context", {"0", "1"});
*graph.add_node() = CreateNode("3", "Exit", {"2"});
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ FrameView frame_view;
+ ASSERT_TRUE(frame_view.InferFromGraph(graph).ok());
+
std::unordered_map<string, std::vector<int>> expected = {
{"0", {}}, {"1", {}}, {"2", {0}}, {"3", {0}}};
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.size(), expected.size());
- std::cout << "Number of frame: " << num_frames << std::endl;
- for (const auto& node : frames) {
- std::cout << node.first->name() << ": ";
- EXPECT_EQ(node.second.size(), expected[node.first->name()].size());
- for (int i = 0; i < node.second.size(); i++) {
- EXPECT_EQ(expected[node.first->name()][i], node.second[i]);
- std::cout << node.second[i] << " ";
- }
- std::cout << std::endl;
+
+ EXPECT_EQ(frame_view.num_frames(), 1);
+ for (const NodeDef& node : graph.node()) {
+ std::vector<int> expected_frames = expected[node.name()];
+ std::vector<int> node_frames = frame_view.Frames(node);
+ EXPECT_EQ(expected_frames, node_frames);
}
}
-TEST_F(IdentifyFramesTest, ExitOutput) {
+TEST_F(FrameViewTest, ExitOutput) {
GraphDef graph;
*graph.add_node() = CreateNode("0", {});
- *graph.add_node() =
- CreateNode("1", "Enter", "map/while/while_context", {"0"});
+ *graph.add_node() = CreateNode("1", "Enter", "while/context", {"0"});
*graph.add_node() = CreateNode("2", "Exit", {"1"});
*graph.add_node() = CreateNode("3", {});
*graph.add_node() = CreateNode("4", {"2", "3"});
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ FrameView frame_view;
+ ASSERT_TRUE(frame_view.InferFromGraph(graph).ok());
+
std::unordered_map<string, std::vector<int>> expected = {
{"0", {}}, {"1", {0}}, {"2", {0}}, {"3", {}}, {"4", {}}};
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.size(), expected.size());
- std::cout << "Number of frame: " << num_frames << std::endl;
- for (const auto& node : frames) {
- std::cout << node.first->name() << ": ";
- EXPECT_EQ(node.second.size(), expected[node.first->name()].size());
- for (int i = 0; i < node.second.size(); i++) {
- EXPECT_EQ(expected[node.first->name()][i], node.second[i]);
- std::cout << node.second[i] << " ";
- }
- std::cout << std::endl;
+
+ EXPECT_EQ(frame_view.num_frames(), 1);
+ for (const NodeDef& node : graph.node()) {
+ std::vector<int> expected_frames = expected[node.name()];
+ std::vector<int> node_frames = frame_view.Frames(node);
+ EXPECT_EQ(expected_frames, node_frames);
}
}
-TEST_F(IdentifyFramesTest, MultipleEnterNodes) {
+TEST_F(FrameViewTest, MultipleEnterNodes) {
GraphDef graph;
*graph.add_node() = CreateNode("0", {});
- string frame = "map/while/while_context";
- *graph.add_node() = CreateNode("1", "Enter", frame, {"0"});
+ *graph.add_node() = CreateNode("1", "Enter", "while/context", {"0"});
*graph.add_node() = CreateNode("2", {"1"});
*graph.add_node() = CreateNode("5", {});
- *graph.add_node() = CreateNode("4", "Enter", frame, {"5"});
+ *graph.add_node() = CreateNode("4", "Enter", "while/context", {"5"});
*graph.add_node() = CreateNode("3", {"4", "2"});
*graph.add_node() = CreateNode("6", "Merge", {"3", "8"});
*graph.add_node() = CreateNode("7", "Switch", {"6"});
*graph.add_node() = CreateNode("8", "NextIteration", {"7"});
*graph.add_node() = CreateNode("9", "Exit", {"7"});
- std::unordered_map<const NodeDef*, std::vector<int>> frames;
- int num_frames;
- EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ FrameView frame_view;
+ ASSERT_TRUE(frame_view.InferFromGraph(graph).ok());
+
std::unordered_map<string, std::vector<int>> expected = {
{"0", {}}, {"1", {0}}, {"2", {0}}, {"3", {0}}, {"4", {0}},
{"5", {}}, {"6", {0}}, {"7", {0}}, {"8", {0}}, {"9", {0}}};
- EXPECT_EQ(num_frames, 1);
- EXPECT_EQ(frames.size(), expected.size());
- std::cout << "Number of frame: " << num_frames << std::endl;
- for (const auto& node : frames) {
- std::cout << node.first->name() << ": ";
- EXPECT_EQ(node.second.size(), expected[node.first->name()].size());
- for (int i = 0; i < node.second.size(); i++) {
- EXPECT_EQ(expected[node.first->name()][i], node.second[i]);
- std::cout << node.second[i] << " ";
- }
- std::cout << std::endl;
+
+ EXPECT_EQ(frame_view.num_frames(), 1);
+ for (const NodeDef& node : graph.node()) {
+ std::vector<int> expected_frames = expected[node.name()];
+ std::vector<int> node_frames = frame_view.Frames(node);
+ EXPECT_EQ(expected_frames, node_frames);
}
}
+TEST_F(FrameViewTest, ConflictingFrames) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode("1", "Enter", "while/context1", {"0"});
+ *graph.add_node() = CreateNode("2", "Enter", "while/context2", {"1"});
+ *graph.add_node() = CreateNode("3", {"1", "2"});
+
+ FrameView frame_view;
+ ASSERT_FALSE(frame_view.InferFromGraph(graph).ok());
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/grappler/utils/functions.cc b/tensorflow/core/grappler/utils/functions.cc
index 57863a7..f2894a9 100644
--- a/tensorflow/core/grappler/utils/functions.cc
+++ b/tensorflow/core/grappler/utils/functions.cc
@@ -347,6 +347,10 @@
fetch.push_back(output_tensor);
}
}
+
+ // It's unsafe to prune side-effectful ops from the graph instantiated from a
+ // function definition (see inlining in function_optimizer.cc).
+ allowed_optimizations().prune_ops_with_side_effects = false;
}
const string& GrapplerFunctionItem::description() const { return description_; }
@@ -561,7 +565,6 @@
inputs.push_back(std::move(input_expansion));
}
- std::vector<string> keep_nodes;
// Add all function nodes to the function body
for (const NodeDef& func_def_node : func.node_def()) {
NodeDef* new_node = function_body.add_node();
@@ -577,11 +580,6 @@
// Register node output range in a function connectivity.
TF_RETURN_IF_ERROR(RegisterFunctionBodyOutputs(*registration, func_def_node,
&connectivity));
-
- // Ops with side effects must be preserved in a function body.
- if (!IsFreeOfSideEffect(func_def_node)) {
- keep_nodes.push_back(func_def_node.name());
- }
}
// Rewrite inputs to use GraphDef format
@@ -612,12 +610,14 @@
outputs.push_back(std::move(output));
}
+ std::vector<string> keep_ops;
bool is_stateful = signature.is_stateful();
*item = GrapplerFunctionItem(
- /*func_name=*/signature.name(), /*description=*/signature.description(),
+ /*func_name=*/signature.name(),
+ /*description=*/signature.description(),
/*func_attr=*/AttrSlice(&func.attr()), std::move(inputs),
- std::move(outputs), std::move(keep_nodes), graph_def_version, is_stateful,
+ std::move(outputs), std::move(keep_ops), graph_def_version, is_stateful,
std::move(function_body));
return Status::OK();
}
diff --git a/tensorflow/core/grappler/utils/functions_test.cc b/tensorflow/core/grappler/utils/functions_test.cc
index 8639dec..5923850 100644
--- a/tensorflow/core/grappler/utils/functions_test.cc
+++ b/tensorflow/core/grappler/utils/functions_test.cc
@@ -599,8 +599,7 @@
EXPECT_EQ(3, item.function_body().node_size());
EXPECT_EQ(1, item.input_size());
EXPECT_EQ(0, item.output_size());
- ASSERT_EQ(1, item.keep_ops.size());
- EXPECT_EQ("update", item.keep_ops[0]);
+ EXPECT_EQ(false, item.allowed_optimizations().prune_ops_with_side_effects);
}
TEST_F(FunctionsTest, MakeFunctionDef) {
diff --git a/tensorflow/core/grappler/utils_test.cc b/tensorflow/core/grappler/utils_test.cc
index e993391..f5ae398 100644
--- a/tensorflow/core/grappler/utils_test.cc
+++ b/tensorflow/core/grappler/utils_test.cc
@@ -464,6 +464,13 @@
Tensor(bfloat16(std::numeric_limits<int>::min())), t);
}
+TEST_F(UtilsTest, TensorIdToString) {
+ EXPECT_EQ("^foo", TensorIdToString({"foo", -1}));
+ EXPECT_EQ("foo", TensorIdToString({"foo", 0}));
+ EXPECT_EQ("foo:1", TensorIdToString({"foo", 1}));
+ EXPECT_EQ("foo:2", TensorIdToString({"foo", 2}));
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 60accc0..e2234c1 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -29,26 +29,26 @@
load(
"//tensorflow:tensorflow.bzl",
+ "cc_header_only_library",
"if_android",
+ "if_not_windows",
+ "tf_cc_binary",
"tf_cc_test",
"tf_cc_test_mkl",
"tf_cc_tests",
- "tf_cc_binary",
"tf_copts",
"tf_cuda_library",
- "tf_opts_nortti_if_android",
"tf_kernel_library",
"tf_mkl_kernel_library",
- "cc_header_only_library",
- "if_not_windows",
+ "tf_opts_nortti_if_android",
)
load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_tests")
load(
"//tensorflow/core:platform/default/build_config.bzl",
- "tf_proto_library",
"tf_kernel_tests_linkstatic",
+ "tf_proto_library",
)
load(
"//tensorflow/core:platform/default/build_config_root.bzl",
@@ -592,6 +592,7 @@
# #endif
cc_library(
name = "eigen_contraction_kernel",
+ srcs = ["eigen_contraction_kernel.cc"],
hdrs = ["eigen_contraction_kernel.h"],
defines = select({
":mkldnn_contraction_kernel": [
@@ -603,7 +604,7 @@
deps = [
"//third_party/eigen3",
] + select({
- ":mkldnn_contraction_kernel": ["//third_party/intel_mkl_dnn:mkldnn_single_threaded"],
+ ":mkldnn_contraction_kernel": ["@mkl_dnn//:mkldnn_single_threaded"],
"//conditions:default": [],
}),
)
@@ -2196,6 +2197,7 @@
":state",
":training_op_helpers",
":variable_ops",
+ "//tensorflow/core:core_cpu_lib",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
@@ -2755,6 +2757,7 @@
":cholesky_grad",
":cholesky_op",
":determinant_op",
+ ":lu_op",
":matrix_exponential_op",
":matrix_inverse_op",
":matrix_logarithm_op",
@@ -2900,6 +2903,19 @@
deps = LINALG_DEPS,
)
+tf_kernel_library(
+ name = "lu_op",
+ prefix = "lu_op",
+ deps = if_cuda([
+ ":cuda_solvers",
+ ":transpose_functor",
+ ]) + [
+ "//third_party/eigen3",
+ "//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ ],
+)
+
cc_library(
name = "linalg_ops_common",
srcs = ["linalg_ops_common.cc"],
@@ -3698,7 +3714,6 @@
":bounds_check",
":conv_2d",
":eigen_contraction_kernel",
- ":fused_batch_norm_util_gpu",
":ops_util",
":pooling_ops",
"//tensorflow/core:framework",
@@ -6741,6 +6756,31 @@
deps = NN_DEPS + mkl_deps() + [":cwise_op"],
)
+tf_cc_test_mkl(
+ name = "mkl_fused_ops_test",
+ size = "small",
+ srcs = ["mkl_fused_ops_test.cc"],
+ linkstatic = 1,
+ deps = [
+ ":conv_ops",
+ ":image",
+ ":mkl_conv_op",
+ ":mkl_tfconv_op",
+ ":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_transpose_op",
srcs = [
diff --git a/tensorflow/core/kernels/control_flow_ops.cc b/tensorflow/core/kernels/control_flow_ops.cc
index 081ef72..36def4a 100644
--- a/tensorflow/core/kernels/control_flow_ops.cc
+++ b/tensorflow/core/kernels/control_flow_ops.cc
@@ -601,6 +601,13 @@
LoopCondOp::~LoopCondOp() = default;
void LoopCondOp::Compute(OpKernelContext* context) {
+ CancellationManager* cm = context->cancellation_manager();
+ if (cm != nullptr) {
+ bool already_cancelled = cm->IsCancelled();
+ OP_REQUIRES(context, !already_cancelled,
+ errors::Cancelled("Loop execution was cancelled."));
+ }
+
context->set_output(0, context->input(0));
}
diff --git a/tensorflow/core/kernels/conv_grad_filter_ops.cc b/tensorflow/core/kernels/conv_grad_filter_ops.cc
index bc30da4..4e3de33 100644
--- a/tensorflow/core/kernels/conv_grad_filter_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_filter_ops.cc
@@ -903,7 +903,7 @@
auto input_ptr = AsDeviceMemory(transformed_input.template flat<T>().data(),
transformed_input.template flat<T>().size());
- static int64 ConvolveBackwardFilterScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveBackwardFilterScratchSize = GetDnnWorkspaceLimit(
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB by default
);
int device_id = stream->parent()->device_ordinal();
@@ -939,8 +939,8 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize,
- ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize,
+ ctx);
ProfileResult profile_result;
bool cudnn_launch_status =
stream
@@ -977,8 +977,7 @@
AutoTuneConvBwdFilter::GetInstance()->Insert(conv_parameters,
algorithm_config);
}
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize,
- ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize, ctx);
bool cudnn_launch_status =
stream
->ThenConvolveBackwardFilterWithAlgorithm(
diff --git a/tensorflow/core/kernels/conv_grad_input_ops.cc b/tensorflow/core/kernels/conv_grad_input_ops.cc
index e06af15..9f983ed 100644
--- a/tensorflow/core/kernels/conv_grad_input_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_input_ops.cc
@@ -951,10 +951,10 @@
AsDeviceMemory(pre_transformed_in_backprop.template flat<T>().data(),
pre_transformed_in_backprop.template flat<T>().size());
- static int64 ConvolveBackwardDataScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveBackwardDataScratchSize = GetDnnWorkspaceLimit(
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB by default
);
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize, ctx);
int device_id = stream->parent()->device_ordinal();
DataType dtype = out_backprop.dtype();
ConvParameters conv_parameters = {
@@ -988,8 +988,8 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
- ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
+ ctx);
ProfileResult profile_result;
bool cudnn_launch_status =
stream
diff --git a/tensorflow/core/kernels/conv_grad_ops_3d.cc b/tensorflow/core/kernels/conv_grad_ops_3d.cc
index e4c49ef..562a9c8 100644
--- a/tensorflow/core/kernels/conv_grad_ops_3d.cc
+++ b/tensorflow/core/kernels/conv_grad_ops_3d.cc
@@ -1333,7 +1333,7 @@
AsDeviceMemory(pre_transformed_in_backprop.template flat<T>().data(),
pre_transformed_in_backprop.template flat<T>().size());
- static int64 ConvolveBackwardDataScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveBackwardDataScratchSize = GetDnnWorkspaceLimit(
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32); // 4GB by default
const int device_id = stream->parent()->device_ordinal();
@@ -1368,8 +1368,8 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
- context);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
+ context);
ProfileResult profile_result;
bool cudnn_launch_status =
stream
@@ -1405,8 +1405,8 @@
AutoTuneConv3dBwdData::GetInstance()->Insert(conv_parameters,
algorithm_config);
}
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
- context);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardDataScratchSize,
+ context);
bool cudnn_launch_status =
stream
->ThenConvolveBackwardDataWithAlgorithm(
@@ -1739,7 +1739,7 @@
AsDeviceMemory(transformed_input.template flat<T>().data(),
transformed_input.template flat<T>().size());
- static int64 ConvolveBackwardFilterScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveBackwardFilterScratchSize = GetDnnWorkspaceLimit(
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32); // 4GB by default
const int device_id = stream->parent()->device_ordinal();
@@ -1774,8 +1774,8 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(
- ConvolveBackwardFilterScratchSize, context);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize,
+ context);
ProfileResult profile_result;
bool cudnn_launch_status =
stream
@@ -1812,8 +1812,8 @@
AutoTuneConv3dBwdFilter::GetInstance()->Insert(conv_parameters,
algorithm_config);
}
- CudnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize,
- context);
+ DnnScratchAllocator scratch_allocator(ConvolveBackwardFilterScratchSize,
+ context);
bool cudnn_launch_status =
stream
->ThenConvolveBackwardFilterWithAlgorithm(
diff --git a/tensorflow/core/kernels/conv_ops.cc b/tensorflow/core/kernels/conv_ops.cc
index 74857fc..dfba157 100644
--- a/tensorflow/core/kernels/conv_ops.cc
+++ b/tensorflow/core/kernels/conv_ops.cc
@@ -521,8 +521,8 @@
template struct LaunchConv2DOp<CPUDevice, double>;
#if GOOGLE_CUDA
-int64 GetCudnnWorkspaceLimit(const string& envvar_in_mb,
- int64 default_value_in_bytes) {
+int64 GetDnnWorkspaceLimit(const string& envvar_in_mb,
+ int64 default_value_in_bytes) {
const char* workspace_limit_in_mb_str = getenv(envvar_in_mb.c_str());
if (workspace_limit_in_mb_str != nullptr &&
strcmp(workspace_limit_in_mb_str, "") != 0) {
@@ -759,7 +759,7 @@
AsDeviceMemory(transformed_output.template flat<T>().data(),
transformed_output.template flat<T>().size());
- static int64 ConvolveScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveScratchSize = GetDnnWorkspaceLimit(
// default value is in bytes despite the name of the environment variable
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB
);
@@ -803,7 +803,7 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
ProfileResult profile_result;
bool cudnn_launch_status =
stream
@@ -841,7 +841,7 @@
AutoTuneConv::GetInstance()->Insert(conv_parameters, algorithm_config);
}
- CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
bool cudnn_launch_status =
stream
->ThenConvolveWithAlgorithm(input_desc, input_ptr, filter_desc,
diff --git a/tensorflow/core/kernels/conv_ops_3d.cc b/tensorflow/core/kernels/conv_ops_3d.cc
index f20ac93..5a59e20 100644
--- a/tensorflow/core/kernels/conv_ops_3d.cc
+++ b/tensorflow/core/kernels/conv_ops_3d.cc
@@ -407,7 +407,7 @@
AsDeviceMemory(transformed_output.template flat<T>().data(),
transformed_output.template flat<T>().size());
- static int64 ConvolveScratchSize = GetCudnnWorkspaceLimit(
+ static int64 ConvolveScratchSize = GetDnnWorkspaceLimit(
"TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32); // 4GB by default
int device_id = stream->parent()->device_ordinal();
@@ -450,7 +450,7 @@
for (auto profile_algorithm : algorithms) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
- CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
ProfileResult profile_result;
bool cudnn_launch_status =
stream
@@ -486,7 +486,7 @@
AutoTuneConv3d::GetInstance()->Insert(conv_parameters, algorithm_config);
}
- CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
+ DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
bool cudnn_launch_status =
stream
->ThenConvolveWithAlgorithm(input_desc, input_ptr, filter_desc,
diff --git a/tensorflow/core/kernels/conv_ops_gpu.h b/tensorflow/core/kernels/conv_ops_gpu.h
index 21d135d..7a67658 100644
--- a/tensorflow/core/kernels/conv_ops_gpu.h
+++ b/tensorflow/core/kernels/conv_ops_gpu.h
@@ -27,19 +27,19 @@
namespace tensorflow {
-// Get the Cudnn workspace limit from the environment variable, which is in MB.
+// Get the Dnn workspace limit from the environment variable, which is in MB.
// Return the workspace memory limit in bytes. If no value is set, return the
// default value.
-int64 GetCudnnWorkspaceLimit(const string& envvar_in_mb,
- int64 default_value_in_bytes);
+int64 GetDnnWorkspaceLimit(const string& envvar_in_mb,
+ int64 default_value_in_bytes);
// A class to provide scratch-space allocator for Stream-Executor Cudnn
// callback. TensorFlow is responsible for releasing the temporary buffers after
// the kernel finishes.
-class CudnnScratchAllocator : public se::ScratchAllocator {
+class DnnScratchAllocator : public se::ScratchAllocator {
public:
- virtual ~CudnnScratchAllocator() {}
- CudnnScratchAllocator(int64 memory_limit, OpKernelContext* context)
+ virtual ~DnnScratchAllocator() {}
+ DnnScratchAllocator(int64 memory_limit, OpKernelContext* context)
: memory_limit_(memory_limit), total_byte_size_(0), context_(context) {}
int64 GetMemoryLimitInBytes(se::Stream* stream) override {
return memory_limit_;
diff --git a/tensorflow/core/kernels/cwise_op_arg.cc b/tensorflow/core/kernels/cwise_op_arg.cc
index 62ffa07..ea659fa 100644
--- a/tensorflow/core/kernels/cwise_op_arg.cc
+++ b/tensorflow/core/kernels/cwise_op_arg.cc
@@ -26,9 +26,7 @@
REGISTER_COMPLEX(CPU, float, complex64);
REGISTER_COMPLEX(CPU, double, complex128);
-// TODO: Enable GPU support for angle op after resolving
-// build failures on GPU (See #10643 for context).
-#if 0 && GOOGLE_CUDA
+#if GOOGLE_CUDA
REGISTER_COMPLEX(GPU, float, complex64);
REGISTER_COMPLEX(GPU, double, complex128);
#endif
diff --git a/tensorflow/core/kernels/cwise_op_gpu_arg.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_arg.cu.cc
index 9b3f820..34028e9 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_arg.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_arg.cu.cc
@@ -13,9 +13,7 @@
limitations under the License.
==============================================================================*/
-// TODO: Enable GPU support for angle op after resolving
-// build failures on GPU (See #10643 for context).
-#if 0 && GOOGLE_CUDA
+#if GOOGLE_CUDA
#include "tensorflow/core/kernels/cwise_ops_gpu_common.cu.h"
diff --git a/tensorflow/core/kernels/cwise_op_squared_difference.cc b/tensorflow/core/kernels/cwise_op_squared_difference.cc
index 78fefc6..d0ff271 100644
--- a/tensorflow/core/kernels/cwise_op_squared_difference.cc
+++ b/tensorflow/core/kernels/cwise_op_squared_difference.cc
@@ -16,8 +16,8 @@
#include "tensorflow/core/kernels/cwise_ops_common.h"
namespace tensorflow {
-REGISTER5(BinaryOp, CPU, "SquaredDifference", functor::squared_difference,
- float, Eigen::half, double, int32, int64);
+REGISTER7(BinaryOp, CPU, "SquaredDifference", functor::squared_difference,
+ float, Eigen::half, double, int32, int64, complex64, complex128);
#if GOOGLE_CUDA
REGISTER4(BinaryOp, GPU, "SquaredDifference", functor::squared_difference,
float, Eigen::half, double, int64);
diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h
index 313def9..a22d767 100644
--- a/tensorflow/core/kernels/cwise_ops.h
+++ b/tensorflow/core/kernels/cwise_ops.h
@@ -29,6 +29,28 @@
namespace Eigen {
namespace internal {
+#if GOOGLE_CUDA
+template <>
+struct scalar_arg_op<std::complex<float>> {
+ EIGEN_EMPTY_STRUCT_CTOR(scalar_arg_op)
+ typedef typename Eigen::NumTraits<std::complex<float>>::Real result_type;
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const float operator()(
+ const std::complex<float>& a) const {
+ return ::atan2f(a.imag(), a.real());
+ }
+};
+
+template <>
+struct scalar_arg_op<std::complex<double>> {
+ EIGEN_EMPTY_STRUCT_CTOR(scalar_arg_op)
+ typedef typename Eigen::NumTraits<std::complex<double>>::Real result_type;
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const double operator()(
+ const std::complex<double>& a) const {
+ return ::atan2(a.imag(), a.real());
+ }
+};
+#endif
+
template <typename T>
struct scalar_asinh_op {
EIGEN_EMPTY_STRUCT_CTOR(scalar_asinh_op)
@@ -296,27 +318,32 @@
}
};
-// Functor that enables composition of multiple Eigen functors.
-template <typename Scalar, typename UnaryFunctor, typename BinaryFunctor>
-struct scalar_compose_op {
+// Functor that enables squared difference functor.
+template <typename Scalar>
+struct scalar_squared_difference_op {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar
operator()(const Scalar& a, const Scalar& b) const {
- return UnaryFunctor()(BinaryFunctor()(a, b));
+ const Scalar v = scalar_difference_op<Scalar>()(a, b);
+ return scalar_product_op<Scalar>()(v, scalar_conjugate_op<Scalar>()(v));
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet
packetOp(const Packet& a, const Packet& b) const {
- return UnaryFunctor().packetOp(BinaryFunctor().packetOp(a, b));
+ const Packet v = scalar_difference_op<Scalar>().packetOp(a, b);
+ return scalar_product_op<Scalar>().packetOp(
+ v, scalar_conjugate_op<Scalar>().packetOp(v));
}
};
-template <typename Scalar, typename UnaryFunctor, typename BinaryFunctor>
-struct functor_traits<scalar_compose_op<Scalar, UnaryFunctor, BinaryFunctor>> {
+template <typename Scalar>
+struct functor_traits<scalar_squared_difference_op<Scalar>> {
enum {
- Cost = functor_traits<UnaryFunctor>::Cost +
- functor_traits<BinaryFunctor>::Cost,
- PacketAccess = functor_traits<UnaryFunctor>::PacketAccess &&
- functor_traits<BinaryFunctor>::PacketAccess
+ Cost = functor_traits<scalar_difference_op<Scalar>>::Cost +
+ functor_traits<scalar_conjugate_op<Scalar>>::Cost +
+ functor_traits<scalar_product_op<Scalar>>::Cost,
+ PacketAccess = functor_traits<scalar_difference_op<Scalar>>::PacketAccess &&
+ functor_traits<scalar_conjugate_op<Scalar>>::PacketAccess &&
+ functor_traits<scalar_product_op<Scalar>>::PacketAccess
};
};
@@ -775,7 +802,7 @@
// pow(x, y) = x ^ y
// maximum(x, y) = x > y ? x : y
// minimum(x, y) = x < y ? x : y
-// squared_difference(x, y) = (x - y) * (x - y)
+// squared_difference(x, y) = conj(x - y) * (x - y)
template <typename T>
struct add : base<T, Eigen::internal::scalar_sum_op<T>> {
@@ -885,9 +912,7 @@
template <typename T>
struct squared_difference
- : base<T, Eigen::internal::scalar_compose_op<
- T, Eigen::internal::scalar_square_op<T>,
- Eigen::internal::scalar_difference_op<T>>> {};
+ : base<T, Eigen::internal::scalar_squared_difference_op<T>> {};
template <typename T>
struct xdivy : base<T, Eigen::internal::xdivy_op<T>> {};
diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc
index d7b231b..f00b38e 100644
--- a/tensorflow/core/kernels/data/cache_dataset_ops.cc
+++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc
@@ -14,6 +14,7 @@
==============================================================================*/
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
+#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/env.h"
@@ -564,9 +565,7 @@
class MemoryDataset : public DatasetBase {
public:
explicit MemoryDataset(OpKernelContext* ctx, const DatasetBase* input)
- : DatasetBase(DatasetContext(ctx)),
- input_(input),
- cache_(new MemoryCache()) {
+ : DatasetBase(DatasetContext(ctx)), input_(input) {
input->Ref();
}
@@ -574,8 +573,8 @@
std::unique_ptr<IteratorBase> MakeIteratorInternal(
const string& prefix) const override {
- return std::unique_ptr<IteratorBase>(new MemoryIterator(
- {this, strings::StrCat(prefix, "::MemoryCache")}, cache_));
+ return std::unique_ptr<IteratorBase>(
+ new MemoryIterator({this, strings::StrCat(prefix, "::MemoryCache")}));
}
const DataTypeVector& output_dtypes() const override {
@@ -611,10 +610,12 @@
// The expected use is that a single `MemoryWriterIterator` populates the
// cache with dataset elements. Once all elements are cached, the cache can
// be used by one or more `MemoryReaderIterator`s.
- class MemoryCache {
+ class MemoryCache : public ResourceBase {
public:
MemoryCache() = default;
+ string DebugString() override { return "CacheDataset::MemoryCache"; }
+
// Marks the cache as completed.
void Complete() {
mutex_lock l(mu_);
@@ -681,15 +682,25 @@
class MemoryIterator : public DatasetIterator<MemoryDataset> {
public:
- explicit MemoryIterator(const Params& params,
- const std::shared_ptr<MemoryCache>& cache)
- : DatasetIterator<MemoryDataset>(params), cache_(cache) {
- mode_ = cache->MaybeClaim() ? Mode::write : Mode::read;
- InitializeIterator();
- }
+ explicit MemoryIterator(const Params& params)
+ : DatasetIterator<MemoryDataset>(params) {}
+
+ ~MemoryIterator() override { cache_->Unref(); }
Status Initialize(IteratorContext* ctx) override {
mutex_lock l(mu_);
+ // Use the resource manager in the iterator context to get / create
+ // a cache.
+ ResourceMgr* mgr = ctx->resource_mgr();
+ const string name =
+ strings::StrCat(prefix(), "::", dataset()->name(), "::MemoryCache");
+ TF_RETURN_IF_ERROR(mgr->LookupOrCreate<MemoryCache>(
+ "tf_data", name, &cache_, [](MemoryCache** cache) {
+ *cache = new MemoryCache();
+ return Status::OK();
+ }));
+ mode_ = cache_->MaybeClaim() ? Mode::write : Mode::read;
+ InitializeIterator();
if (mode_ == Mode::read && !cache_->IsCompleted()) {
return errors::Internal(
"Cache should only be read after it has been completed.");
@@ -788,8 +799,7 @@
private:
class MemoryWriterIterator : public DatasetIterator<MemoryDataset> {
public:
- explicit MemoryWriterIterator(const Params& params,
- const std::shared_ptr<MemoryCache>& cache)
+ explicit MemoryWriterIterator(const Params& params, MemoryCache* cache)
: DatasetIterator<MemoryDataset>(params), cache_(cache) {
CHECK(cache_);
}
@@ -822,6 +832,7 @@
cache_->Complete();
return Status::OK();
}
+ RecordBufferEnqueue(ctx, *out_tensors);
cache_->emplace_back(*out_tensors);
return Status::OK();
}
@@ -847,17 +858,46 @@
private:
mutex mu_;
std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
- std::shared_ptr<MemoryCache> cache_;
+ MemoryCache* const cache_ GUARDED_BY(mu_); // not owned.
}; // MemoryWriterIterator
class MemoryReaderIterator : public DatasetIterator<MemoryDataset> {
public:
- explicit MemoryReaderIterator(const Params& params,
- const std::shared_ptr<MemoryCache>& cache)
+ explicit MemoryReaderIterator(const Params& params, MemoryCache* cache)
: DatasetIterator<MemoryDataset>(params), cache_(cache), index_(0) {
CHECK(cache);
}
+ Status Initialize(IteratorContext* ctx) override {
+ // The memory allocated for the cache is owned by the parent
+ // dataset but performance modeling uses the iterator abstraction and
+ // thus we record the memory allocated for the cache here. The caveat
+ // is that this is incorrect if there are concurrent instances of this
+ // iterator.
+ tf_shared_lock l(mu_);
+ for (size_t i = 0; i < cache_->size(); ++i) {
+ RecordBufferEnqueue(ctx, cache_->at(i));
+ }
+ return Status::OK();
+ }
+
+ Status GetNextInternal(IteratorContext* ctx,
+ std::vector<Tensor>* out_tensors,
+ bool* end_of_sequence) override {
+ mutex_lock l(mu_);
+ if (index_ < cache_->size()) {
+ const std::vector<Tensor>& cache_tensors = cache_->at(index_);
+ out_tensors->insert(out_tensors->begin(), cache_tensors.begin(),
+ cache_tensors.end());
+ index_++;
+ *end_of_sequence = false;
+ return Status::OK();
+ } else {
+ *end_of_sequence = true;
+ return Status::OK();
+ }
+ }
+
protected:
std::shared_ptr<model::Node> CreateNode(
IteratorContext* ctx, model::Node::Args args) const override {
@@ -882,26 +922,9 @@
return Status::OK();
}
- Status GetNextInternal(IteratorContext* ctx,
- std::vector<Tensor>* out_tensors,
- bool* end_of_sequence) override {
- mutex_lock l(mu_);
- if (index_ < cache_->size()) {
- const std::vector<Tensor>& cache_tensors = cache_->at(index_);
- out_tensors->insert(out_tensors->begin(), cache_tensors.begin(),
- cache_tensors.end());
- index_++;
- *end_of_sequence = false;
- return Status::OK();
- } else {
- *end_of_sequence = true;
- return Status::OK();
- }
- }
-
private:
mutex mu_;
- const std::shared_ptr<MemoryCache> cache_;
+ MemoryCache* const cache_ GUARDED_BY(mu_); // not owned.
size_t index_ GUARDED_BY(mu_);
}; // MemoryReaderIterator
@@ -918,14 +941,13 @@
}
mutex mu_;
- std::shared_ptr<MemoryCache> cache_;
+ MemoryCache* cache_ GUARDED_BY(mu_); // not owned.
enum Mode { read, write };
Mode mode_ GUARDED_BY(mu_);
std::unique_ptr<IteratorBase> iterator_ GUARDED_BY(mu_);
}; // MemoryIterator
const DatasetBase* const input_;
- const std::shared_ptr<MemoryCache> cache_;
}; // MemoryDataset
}; // CacheDatasetOp
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 df96ccc..3ff3135 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
@@ -51,11 +51,12 @@
std::shared_ptr<std::vector<Tensor>>, StatusCallback)>;
explicit MapAndBatchDatasetOp(OpKernelConstruction* ctx)
- : UnaryDatasetOpKernel(ctx),
- op_version_(ctx->def().op() == "MapAndBatchDataset" ? 1 : 2) {
+ : UnaryDatasetOpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
+ OP_REQUIRES_OK(
+ ctx, ctx->GetAttr("preserve_cardinality", &preserve_cardinality_));
}
protected:
@@ -68,29 +69,12 @@
errors::InvalidArgument("batch_size must be greater than zero."));
int64 num_parallel_calls;
- switch (op_version_) {
- case 1:
- int64 num_parallel_batches;
- OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_batches",
- &num_parallel_batches));
- num_parallel_calls = num_parallel_batches * batch_size;
- OP_REQUIRES(ctx, num_parallel_batches > 0,
- errors::InvalidArgument(
- "num_parallel_batches must be greater than zero."));
- break;
- case 2:
- OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls",
- &num_parallel_calls));
- OP_REQUIRES(ctx,
- num_parallel_calls > 0 || num_parallel_calls == kAutoTune,
- errors::InvalidArgument(
- "num_parallel_calls must be greater than zero."));
- break;
- default:
- OP_REQUIRES(ctx, false,
- errors::Unimplemented("Unsupported operation version %d.",
- op_version_));
- }
+ OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls",
+ &num_parallel_calls));
+ OP_REQUIRES(
+ ctx, num_parallel_calls > 0 || num_parallel_calls == model::kAutoTune,
+ errors::InvalidArgument(
+ "num_parallel_calls must be greater than zero."));
bool drop_remainder;
OP_REQUIRES_OK(ctx,
@@ -146,7 +130,7 @@
*output = new Dataset(ctx, input, func_, batch_size, num_parallel_calls,
drop_remainder, output_types_, output_shapes_,
std::move(captured_func), &ctx->eigen_cpu_device(),
- std::move(map_func));
+ std::move(map_func), preserve_cardinality_);
}
private:
@@ -159,7 +143,7 @@
const std::vector<PartialTensorShape>& output_shapes,
std::unique_ptr<CapturedFunction> captured_func,
const Eigen::ThreadPoolDevice* device,
- MapAndBatchIteratorFunction map_func)
+ MapAndBatchIteratorFunction map_func, bool preserve_cardinality)
: DatasetBase(DatasetContext(ctx)),
input_(input),
func_(func),
@@ -170,7 +154,8 @@
output_shapes_(output_shapes),
captured_func_(std::move(captured_func)),
device_(device),
- map_func_(std::move(map_func)) {
+ map_func_(std::move(map_func)),
+ preserve_cardinality_(preserve_cardinality) {
input_->Ref();
}
@@ -195,8 +180,6 @@
return "MapAndBatchDatasetOp::Dataset";
}
- // TODO(b/120482302): Note that this is inaccurate until MapDataset is
- // modified to preserve cardinality.
int64 Cardinality() const override {
int64 n = input_->Cardinality();
if (n == kInfiniteCardinality || n == kUnknownCardinality) {
@@ -235,6 +218,8 @@
b->BuildAttrValue(func_, &f);
AttrValue other_arguments_types_attr;
b->BuildAttrValue(other_arguments_types, &other_arguments_types_attr);
+ AttrValue preserve_cardinality_attr;
+ b->BuildAttrValue(preserve_cardinality_, &preserve_cardinality_attr);
TF_RETURN_IF_ERROR(b->AddDataset(
this,
@@ -244,7 +229,9 @@
std::make_pair(4, drop_remainder_node)}, // Single tensor inputs.
{std::make_pair(1, other_arguments)}, // Tensor list inputs.
{std::make_pair("f", f),
- std::make_pair("Targuments", other_arguments_types_attr)}, // Attrs
+ std::make_pair("Targuments", other_arguments_types_attr),
+ std::make_pair("preserve_cardinality",
+ preserve_cardinality_attr)}, // Attrs
output));
return Status::OK();
}
@@ -282,9 +269,8 @@
Status Initialize(IteratorContext* ctx) override {
mutex_lock l(*mu_);
- if (num_parallel_calls_->value == kAutoTune) {
+ if (num_parallel_calls_->value == model::kAutoTune) {
num_parallel_calls_->value = ctx->runner_threadpool_size();
- num_parallel_calls_->tunable = true;
}
TF_RETURN_IF_ERROR(
dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_));
@@ -434,6 +420,15 @@
std::shared_ptr<std::vector<Tensor>> return_values =
std::make_shared<std::vector<Tensor>>();
auto done = [this, ctx, result, return_values, offset](Status status) {
+ if (dataset()->preserve_cardinality_ &&
+ errors::IsOutOfRange(status)) {
+ // To guarantee that the transformation preserves the cardinality of
+ // the dataset, we convert `OutOfRange` to `InvalidArgument` as the
+ // former may be interpreted by a caller as the end of sequence.
+ status = errors::InvalidArgument(
+ "Function invocation produced OutOfRangeError: ",
+ status.error_message());
+ }
result->UpdateStatus(status, offset);
if (status.ok()) {
Status allocate_status =
@@ -460,8 +455,8 @@
// TODO(mrry): Add a version of DoParallelConcat that allows us
// to move `tensor` where possible, to speed up string tensor
// batching.
- Status copy_status = ::tensorflow::functor::DoParallelConcat(
- *dataset()->device_, tensor, offset, batch);
+ Status copy_status =
+ batch_util::CopyElementToSlice(tensor, batch, offset);
if (!copy_status.ok()) {
result->UpdateStatus(copy_status, offset);
break;
@@ -546,11 +541,14 @@
bool* end_of_sequence) {
mutex_lock l(result->mu);
if (result->num_elements == 0) {
- *end_of_sequence = true;
- return Status::OK();
+ if (result->status.ok() || errors::IsOutOfRange(result->status)) {
+ *end_of_sequence = true;
+ return Status::OK();
+ } else {
+ *end_of_sequence = false;
+ return result->status;
+ }
}
- // `f` may deliberately raise `errors::OutOfRange` to indicate that we
- // should terminate the iteration early.
if (!result->status.ok() && !errors::IsOutOfRange(result->status)) {
// Deallocate tensors allocated for the output.
result->output.clear();
@@ -580,7 +578,7 @@
} else {
*out_tensors = std::move(result->output);
}
- *end_of_sequence = result->num_elements == 0;
+ *end_of_sequence = false;
return Status::OK();
}
@@ -813,12 +811,13 @@
const std::unique_ptr<CapturedFunction> captured_func_;
const Eigen::ThreadPoolDevice* device_; // not owned
const MapAndBatchIteratorFunction map_func_;
+ const bool preserve_cardinality_;
};
- const int op_version_;
DataTypeVector output_types_;
std::vector<PartialTensorShape> output_shapes_;
NameAttrList func_;
+ bool preserve_cardinality_;
};
REGISTER_KERNEL_BUILDER(
diff --git a/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc
index 7051012..921f8ad 100644
--- a/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/numa_map_and_batch_dataset_op.cc
@@ -59,6 +59,9 @@
OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &func_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
+ // TODO(saeta): Implement support for preserve_cardinality logic.
+ OP_REQUIRES_OK(
+ ctx, ctx->GetAttr("preserve_cardinality", &preserve_cardinality_));
}
protected:
@@ -73,9 +76,10 @@
int64 num_parallel_calls;
OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls",
&num_parallel_calls));
- OP_REQUIRES(ctx, num_parallel_calls > 0 || num_parallel_calls == kAutoTune,
- errors::InvalidArgument(
- "num_parallel_calls must be greater than zero."));
+ OP_REQUIRES(
+ ctx, num_parallel_calls > 0 || num_parallel_calls == model::kAutoTune,
+ errors::InvalidArgument(
+ "num_parallel_calls must be greater than zero."));
bool drop_remainder;
OP_REQUIRES_OK(ctx,
@@ -133,6 +137,8 @@
return "NumaMapAndBatchDatasetOp::Dataset";
}
+ // TODO(b/120482302): Note that this is inaccurate until
+ // NumaMapAndBatchMapDataset modified to preserve cardinality.
int64 Cardinality() const override {
int64 n = input_->Cardinality();
if (n == kInfiniteCardinality || n == kUnknownCardinality) {
@@ -209,9 +215,8 @@
Status Initialize(IteratorContext* ctx) override {
mutex_lock l(*mu_);
- if (num_parallel_calls_->value == kAutoTune) {
+ if (num_parallel_calls_->value == model::kAutoTune) {
num_parallel_calls_->value = ctx->runner_threadpool_size();
- num_parallel_calls_->tunable = true;
}
TF_RETURN_IF_ERROR(
dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_));
@@ -1138,6 +1143,7 @@
DataTypeVector output_types_;
std::vector<PartialTensorShape> output_shapes_;
NameAttrList func_;
+ bool preserve_cardinality_;
};
REGISTER_KERNEL_BUILDER(
diff --git a/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc b/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
index 40771a5..ea99a8b 100644
--- a/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
@@ -187,7 +187,8 @@
new ParseExampleFunctor(this));
return NewParallelMapIterator(
{this, strings::StrCat(prefix, "::ParseExample")}, input_,
- std::move(parse_example_functor), num_parallel_calls_, sloppy_);
+ std::move(parse_example_functor), num_parallel_calls_, sloppy_,
+ /*preserve_cardinality=*/true);
}
const DataTypeVector& output_dtypes() const override {
diff --git a/tensorflow/core/kernels/data/experimental/scan_dataset_op.cc b/tensorflow/core/kernels/data/experimental/scan_dataset_op.cc
index 2eceab1..0d9a629 100644
--- a/tensorflow/core/kernels/data/experimental/scan_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/scan_dataset_op.cc
@@ -37,6 +37,8 @@
OP_REQUIRES_OK(ctx, ctx->GetAttr("Tstate", &state_types_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
+ OP_REQUIRES_OK(
+ ctx, ctx->GetAttr("preserve_cardinality", &preserve_cardinality_));
}
void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
@@ -53,7 +55,7 @@
*output = new Dataset(ctx, input, func_, std::move(initial_state),
std::move(captured_func), state_types_, output_types_,
- output_shapes_);
+ output_shapes_, preserve_cardinality_);
}
private:
@@ -64,7 +66,8 @@
std::unique_ptr<CapturedFunction> captured_func,
const DataTypeVector& state_types,
const DataTypeVector& output_types,
- const std::vector<PartialTensorShape>& output_shapes)
+ const std::vector<PartialTensorShape>& output_shapes,
+ bool preserve_cardinality)
: DatasetBase(DatasetContext(ctx)),
input_(input),
func_(func),
@@ -72,7 +75,8 @@
captured_func_(std::move(captured_func)),
state_types_(state_types),
output_types_(output_types),
- output_shapes_(output_shapes) {
+ output_shapes_(output_shapes),
+ preserve_cardinality_(preserve_cardinality) {
input_->Ref();
}
@@ -93,8 +97,6 @@
string DebugString() const override { return "ScanDatasetOp::Dataset"; }
- // TODO(b/120482302): Note that this is inaccurate until MapDataset is
- // modified to preserve cardinality.
int64 Cardinality() const override { return input_->Cardinality(); }
protected:
@@ -127,12 +129,15 @@
b->BuildAttrValue(state_types_, &state_types);
AttrValue other_arguments_types_attr;
b->BuildAttrValue(other_arguments_types, &other_arguments_types_attr);
+ AttrValue preserve_cardinality_attr;
+ b->BuildAttrValue(preserve_cardinality_, &preserve_cardinality_attr);
TF_RETURN_IF_ERROR(
b->AddDataset(this, {{0, input_node}},
{{1, initial_state_nodes}, {2, other_arguments}},
{{"f", f},
{"Tstate", state_types},
- {"Targuments", other_arguments_types_attr}},
+ {"Targuments", other_arguments_types_attr},
+ {"preserve_cardinality", preserve_cardinality_attr}},
output));
return Status::OK();
}
@@ -207,10 +212,19 @@
out_tensors->push_back(std::move(state_and_output[i]));
}
} else if (errors::IsOutOfRange(s)) {
- // `f` may deliberately raise `errors::OutOfRange` to indicate
- // that we should terminate the iteration early.
- *end_of_sequence = true;
- return Status::OK();
+ if (dataset()->preserve_cardinality_) {
+ // To guarantee that the transformation preserves the cardinality of
+ // the dataset, we convert `OutOfRange` to `InvalidArgument` as the
+ // former may be interpreted by a caller as the end of sequence.
+ return errors::InvalidArgument(
+ "Function invocation produced OutOfRangeError: ",
+ s.error_message());
+ } else {
+ // `f` may deliberately raise `errors::OutOfRange` to indicate
+ // that we should terminate the iteration early.
+ *end_of_sequence = true;
+ return Status::OK();
+ }
}
return s;
}
@@ -267,12 +281,14 @@
const DataTypeVector state_types_;
const DataTypeVector output_types_;
const std::vector<PartialTensorShape> output_shapes_;
+ const bool preserve_cardinality_;
};
DataTypeVector state_types_;
DataTypeVector output_types_;
std::vector<PartialTensorShape> output_shapes_;
NameAttrList func_;
+ bool preserve_cardinality_;
};
REGISTER_KERNEL_BUILDER(Name("ExperimentalScanDataset").Device(DEVICE_CPU),
diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index cb7477f..d5b4bfa 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -77,6 +77,7 @@
params.lib = captured_state->lib;
params.function_handle_cache =
captured_state->function_handle_cache.get();
+ params.resource_mgr = &captured_state->resource_mgr;
return captured_state->iterator->GetNext(
IteratorContext(std::move(params)), out_tensors, end_of_sequence);
} else {
@@ -135,8 +136,8 @@
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(nullptr);
TF_RETURN_IF_ERROR(ctx->function_library()->Clone(&flib_def, &pflr, &lib));
TF_RETURN_IF_ERROR(flib_def->AddLibrary(graph_def.library()));
- std::unique_ptr<State> new_state(
- new State(std::move(flib_def), std::move(pflr), lib, nullptr));
+ std::unique_ptr<State> new_state(new State(
+ std::move(flib_def), std::move(pflr), lib, nullptr /* iterator */));
TF_RETURN_IF_ERROR(
graph_runner.Run(&graph, new_state->lib, {}, {output_node}, &outputs));
@@ -145,6 +146,7 @@
IteratorContext::Params params(ctx);
params.lib = new_state->lib;
params.function_handle_cache = new_state->function_handle_cache.get();
+ params.resource_mgr = &new_state->resource_mgr;
TF_RETURN_IF_ERROR(dataset->MakeIterator(IteratorContext(std::move(params)),
"Iterator", &new_state->iterator));
TF_RETURN_IF_ERROR(
@@ -156,6 +158,7 @@
IteratorContext::Params params(ctx);
params.lib = new_state->lib;
params.function_handle_cache = new_state->function_handle_cache.get();
+ params.resource_mgr = &new_state->resource_mgr;
DeviceBase* device = new_state->lib->device();
params.allocator_getter = [device](AllocatorAttributes attrs) {
return device->GetAllocator(attrs);
@@ -180,7 +183,8 @@
tf_shared_lock l(mu_);
new_state.reset(new State(iterator_state_->flib_def,
iterator_state_->pflr, iterator_state_->lib,
- nullptr, nullptr));
+ nullptr /* function_handle_cache */,
+ nullptr /* iterator */));
}
// Ensure that the iterator has access to all functions in the current
@@ -212,6 +216,7 @@
IteratorContext::Params params(ctx);
params.lib = new_state->lib;
params.function_handle_cache = new_state->function_handle_cache.get();
+ params.resource_mgr = &new_state->resource_mgr;
TF_RETURN_IF_ERROR(dataset->MakeIterator(IteratorContext(std::move(params)),
"Iterator", &iterator));
TF_RETURN_IF_ERROR(
@@ -259,6 +264,7 @@
std::shared_ptr<ProcessFunctionLibraryRuntime> pflr;
FunctionLibraryRuntime* lib = nullptr; // not owned.
std::unique_ptr<FunctionHandleCache> function_handle_cache;
+ ResourceMgr resource_mgr;
std::unique_ptr<IteratorBase> iterator;
};
diff --git a/tensorflow/core/kernels/data/map_dataset_op.cc b/tensorflow/core/kernels/data/map_dataset_op.cc
index e7ec02f..fc6e93a 100644
--- a/tensorflow/core/kernels/data/map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/map_dataset_op.cc
@@ -40,6 +40,8 @@
OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("use_inter_op_parallelism",
&use_inter_op_parallelism_));
+ OP_REQUIRES_OK(
+ ctx, ctx->GetAttr("preserve_cardinality", &preserve_cardinality_));
}
void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
@@ -86,9 +88,10 @@
};
}
- *output = new Dataset(ctx, input, func_, std::move(captured_func),
- output_types_, output_shapes_,
- use_inter_op_parallelism_, std::move(map_func));
+ *output =
+ new Dataset(ctx, input, func_, std::move(captured_func), output_types_,
+ output_shapes_, use_inter_op_parallelism_,
+ std::move(map_func), preserve_cardinality_);
}
private:
@@ -99,11 +102,13 @@
std::unique_ptr<CapturedFunction> captured_func,
const DataTypeVector& output_types,
const std::vector<PartialTensorShape>& output_shapes,
- bool use_inter_op_parallelism, MapIteratorFunction map_func)
+ bool use_inter_op_parallelism, MapIteratorFunction map_func,
+ bool preserve_cardinality)
: DatasetBase(DatasetContext(ctx)),
input_(input),
func_(func),
use_inter_op_parallelism_(use_inter_op_parallelism),
+ preserve_cardinality_(preserve_cardinality),
captured_func_(std::move(captured_func)),
output_types_(output_types),
output_shapes_(output_shapes),
@@ -128,8 +133,6 @@
string DebugString() const override { return "MapDatasetOp::Dataset"; }
- // TODO(b/120482302): Note that this is inaccurate until MapDataset is
- // modified to preserve cardinality.
int64 Cardinality() const override { return input_->Cardinality(); }
protected:
@@ -165,13 +168,19 @@
b->BuildAttrValue(use_inter_op_parallelism_,
&use_inter_op_parallelism_attr);
+ // Attr: preserve_cardinality
+ AttrValue preserve_cardinality_attr;
+ b->BuildAttrValue(preserve_cardinality_, &preserve_cardinality_attr);
+
TF_RETURN_IF_ERROR(b->AddDataset(
this, {std::make_pair(0, input_graph_node)}, // Single tensor inputs.
{std::make_pair(1, other_arguments)}, // Tensor list inputs.
{std::make_pair("f", f_attr),
std::make_pair("Targuments", other_arguments_types_attr),
std::make_pair("use_inter_op_parallelism",
- use_inter_op_parallelism_attr)}, // Attrs
+ use_inter_op_parallelism_attr),
+ std::make_pair("preserve_cardinality",
+ preserve_cardinality_attr)}, // Attrs
output));
return Status::OK();
}
@@ -206,10 +215,19 @@
Status s = map_func_(ctx, instantiated_captured_func_.get(), args,
out_tensors);
if (errors::IsOutOfRange(s)) {
- // `f` may deliberately raise `errors::OutOfRange` to indicate
- // that we should terminate the iteration early.
- *end_of_sequence = true;
- return Status::OK();
+ if (dataset()->preserve_cardinality_) {
+ // To guarantee that the transformation preserves the cardinality of
+ // the dataset, we convert `OutOfRange` to `InvalidArgument` as the
+ // former may be interpreted by a caller as the end of sequence.
+ return errors::InvalidArgument(
+ "Function invocation produced OutOfRangeError: ",
+ s.error_message());
+ } else {
+ // `f` may deliberately raise `errors::OutOfRange` to indicate
+ // that we should terminate the iteration early.
+ *end_of_sequence = true;
+ return Status::OK();
+ }
} else {
return s;
}
@@ -242,6 +260,7 @@
const DatasetBase* const input_;
const NameAttrList func_;
const bool use_inter_op_parallelism_;
+ const bool preserve_cardinality_;
const std::unique_ptr<CapturedFunction> captured_func_;
const DataTypeVector output_types_;
const std::vector<PartialTensorShape> output_shapes_;
@@ -252,6 +271,7 @@
std::vector<PartialTensorShape> output_shapes_;
NameAttrList func_;
bool use_inter_op_parallelism_;
+ bool preserve_cardinality_;
};
REGISTER_KERNEL_BUILDER(Name("MapDataset").Device(DEVICE_CPU), MapDatasetOp);
diff --git a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
index a070456..ba2125a 100644
--- a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
+++ b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
@@ -98,6 +98,7 @@
IteratorContext::Params params(ctx);
params.lib = lib_;
params.function_handle_cache = function_handle_cache_.get();
+ params.resource_mgr = &resource_mgr_;
IteratorContext iter_ctx(std::move(params));
tf_shared_lock l(mu_);
multi_device_buffer_->GetNextFromShard(
@@ -125,6 +126,8 @@
return function_handle_cache_.get();
}
+ ResourceMgr* resource_mgr() { return &resource_mgr_; }
+
private:
// A private class that uses a background thread to keep a per device buffer
// full.
@@ -350,6 +353,7 @@
const std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
FunctionLibraryRuntime* const lib_ = nullptr; // not owned.
const std::unique_ptr<FunctionHandleCache> function_handle_cache_;
+ ResourceMgr resource_mgr_;
std::shared_ptr<const FunctionLibraryDefinition> lib_def_ GUARDED_BY(mu_);
int64 incarnation_id_ GUARDED_BY(mu_) = 0;
@@ -477,6 +481,7 @@
IteratorContext::Params params(ctx);
params.lib = resource->lib();
params.function_handle_cache = resource->function_handle_cache();
+ params.resource_mgr = resource->resource_mgr();
IteratorContext iter_ctx(std::move(params));
OP_REQUIRES_OK(
ctx, dataset->MakeIterator(std::move(iter_ctx), "Iterator", &iterator));
diff --git a/tensorflow/core/kernels/data/optimize_dataset_op.cc b/tensorflow/core/kernels/data/optimize_dataset_op.cc
index ab184c2..9c50d80 100644
--- a/tensorflow/core/kernels/data/optimize_dataset_op.cc
+++ b/tensorflow/core/kernels/data/optimize_dataset_op.cc
@@ -305,8 +305,9 @@
// removing unused graph nodes)
// TODO(b/118175421): This should be part of the tf.data optimization
// pass manager.
- for (const auto& optimizer : {"pruning", "function", "constfold",
- "shape", "arithmetic", "dependency"}) {
+ // TODO(b/120437209): Apply `constfold` optimization when it is fixed.
+ for (const auto& optimizer :
+ {"pruning", "function", "shape", "arithmetic", "dependency"}) {
rewriter_config.add_optimizers(optimizer);
}
}
diff --git a/tensorflow/core/kernels/data/optional_ops.cc b/tensorflow/core/kernels/data/optional_ops.cc
index d8a7f21..a406f74 100644
--- a/tensorflow/core/kernels/data/optional_ops.cc
+++ b/tensorflow/core/kernels/data/optional_ops.cc
@@ -159,9 +159,13 @@
to_values.reserve(from_values.size());
for (const Tensor& t : from_values) {
if (DMAHelper::CanUseDMA(&t) || t.dtype() == DT_VARIANT) {
- Tensor tmp(t.dtype());
- TF_RETURN_IF_ERROR(copy(t, &tmp));
- to_values.push_back(std::move(tmp));
+ // NOTE(skyewm): we're careful to make sure the lifetime of the 'to'
+ // Tensor passed to `copy` (i.e. to_values.back()) is the same as the
+ // returned 'to' OptionalVariant. This is because `copy` may spawn async
+ // callbacks that don't run until after this function returns and access
+ // the 'to' Tensor (e.g. BaseGPUDevice::MaybeCopyTensorToGPU).
+ to_values.emplace_back(t.dtype());
+ TF_RETURN_IF_ERROR(copy(t, &to_values.back()));
} else {
to_values.push_back(t);
}
diff --git a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
index 5d33d1e..f844a00 100644
--- a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
@@ -76,9 +76,10 @@
int64 num_parallel_calls;
OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls",
&num_parallel_calls));
- OP_REQUIRES(ctx, num_parallel_calls > 0 || num_parallel_calls == kAutoTune,
- errors::InvalidArgument(
- "num_parallel_calls must be greater than zero."));
+ OP_REQUIRES(
+ ctx, num_parallel_calls > 0 || num_parallel_calls == model::kAutoTune,
+ errors::InvalidArgument(
+ "num_parallel_calls must be greater than zero."));
OP_REQUIRES(
ctx, num_parallel_calls <= cycle_length,
errors::InvalidArgument(
@@ -220,9 +221,8 @@
Status Initialize(IteratorContext* ctx) override {
mutex_lock l(*mu_);
- if (num_parallel_calls_->value == kAutoTune) {
+ if (num_parallel_calls_->value == model::kAutoTune) {
num_parallel_calls_->value = dataset()->cycle_length_;
- num_parallel_calls_->tunable = true;
}
TF_RETURN_IF_ERROR(
dataset()->input_->MakeIterator(ctx, prefix(), &input_impl_));
@@ -256,6 +256,7 @@
if (result->status.ok()) {
*out_tensors = std::move(result->return_values);
+ RecordBufferDequeue(ctx, *out_tensors);
}
*end_of_sequence = false;
return result->status;
@@ -394,6 +395,7 @@
if (end_of_input) {
result->skip = true;
}
+ RecordBufferEnqueue(ctx.get(), result->return_values);
{
mutex_lock l(*mu_);
result->notification.Notify();
diff --git a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
index 7c589a6..5c09b2d 100644
--- a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
@@ -41,6 +41,8 @@
OP_REQUIRES_OK(ctx, ctx->GetAttr("use_inter_op_parallelism",
&use_inter_op_parallelism_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("sloppy", &sloppy_));
+ OP_REQUIRES_OK(
+ ctx, ctx->GetAttr("preserve_cardinality", &preserve_cardinality_));
}
protected:
@@ -49,9 +51,10 @@
int32 num_parallel_calls;
OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls",
&num_parallel_calls));
- OP_REQUIRES(ctx, num_parallel_calls > 0 || num_parallel_calls == kAutoTune,
- errors::InvalidArgument(
- "num_parallel_calls must be greater than zero."));
+ OP_REQUIRES(
+ ctx, num_parallel_calls > 0 || num_parallel_calls == model::kAutoTune,
+ errors::InvalidArgument(
+ "num_parallel_calls must be greater than zero."));
std::unique_ptr<CapturedFunction> captured_func;
OP_REQUIRES_OK(ctx, CapturedFunction::Create(func_, ctx, "other_arguments",
@@ -61,9 +64,10 @@
std::vector<int> indices;
OP_REQUIRES_OK(ctx, ComputeShortCircuitIndices(ctx, func_, &indices));
- *output = new Dataset(ctx, input, func_, num_parallel_calls, output_types_,
- output_shapes_, use_inter_op_parallelism_, sloppy_,
- std::move(captured_func), indices);
+ *output =
+ new Dataset(ctx, input, func_, num_parallel_calls, output_types_,
+ output_shapes_, use_inter_op_parallelism_, sloppy_,
+ std::move(captured_func), indices, preserve_cardinality_);
}
private:
@@ -75,7 +79,7 @@
const std::vector<PartialTensorShape>& output_shapes,
bool use_inter_op_parallelism, bool sloppy,
std::unique_ptr<CapturedFunction> captured_func,
- const std::vector<int> indices)
+ const std::vector<int> indices, bool preserve_cardinality)
: DatasetBase(DatasetContext(ctx)),
input_(input),
func_(func),
@@ -84,6 +88,7 @@
output_shapes_(output_shapes),
use_inter_op_parallelism_(use_inter_op_parallelism),
sloppy_(sloppy),
+ preserve_cardinality_(preserve_cardinality),
captured_func_(std::move(captured_func)),
indices_(indices),
can_move_(indices.empty() ? std::vector<bool>()
@@ -103,7 +108,8 @@
}
return NewParallelMapIterator(
{this, strings::StrCat(prefix, "::ParallelMap")}, input_,
- std::move(parallel_map_functor), num_parallel_calls_, sloppy_);
+ std::move(parallel_map_functor), num_parallel_calls_, sloppy_,
+ preserve_cardinality_);
}
const DataTypeVector& output_dtypes() const override {
@@ -118,8 +124,6 @@
return "ParallelMapDatasetOp::Dataset";
}
- // TODO(b/120482302): Note that this is inaccurate until MapDataset is
- // modified to preserve cardinality.
int64 Cardinality() const override { return input_->Cardinality(); }
protected:
@@ -165,6 +169,10 @@
AttrValue sloppy_attr;
b->BuildAttrValue(sloppy_, &sloppy_attr);
+ // Attr: preserve_cardinality
+ AttrValue preserve_cardinality_attr;
+ b->BuildAttrValue(preserve_cardinality_, &preserve_cardinality_attr);
+
TF_RETURN_IF_ERROR(b->AddDataset(
this,
{std::make_pair(0, input_graph_node),
@@ -174,7 +182,9 @@
std::make_pair("Targuments", other_arguments_types_attr),
std::make_pair("use_inter_op_parallelism",
use_inter_op_parallelism_attr),
- std::make_pair("sloppy", sloppy_attr)}, // Attrs
+ std::make_pair("sloppy", sloppy_attr),
+ std::make_pair("preserve_cardinality",
+ preserve_cardinality_attr)}, // Attrs
output));
return Status::OK();
}
@@ -252,6 +262,7 @@
const std::vector<PartialTensorShape> output_shapes_;
const bool use_inter_op_parallelism_;
const bool sloppy_;
+ const bool preserve_cardinality_;
const std::unique_ptr<CapturedFunction> captured_func_;
const std::vector<int> indices_;
const std::vector<bool> can_move_;
@@ -261,6 +272,7 @@
std::vector<PartialTensorShape> output_shapes_;
bool use_inter_op_parallelism_;
bool sloppy_;
+ bool preserve_cardinality_;
NameAttrList func_;
};
diff --git a/tensorflow/core/kernels/data/parallel_map_iterator.cc b/tensorflow/core/kernels/data/parallel_map_iterator.cc
index 02ccf6b..b62e705 100644
--- a/tensorflow/core/kernels/data/parallel_map_iterator.cc
+++ b/tensorflow/core/kernels/data/parallel_map_iterator.cc
@@ -32,20 +32,34 @@
class ParallelMapIterator : public DatasetBaseIterator {
public:
- ParallelMapIterator(const typename DatasetBaseIterator::BaseParams& params,
- const DatasetBase* input_dataset,
- std::unique_ptr<ParallelMapFunctor> parallel_map_functor,
- int32 num_parallel_calls, bool sloppy)
- : DatasetBaseIterator(params),
+ struct Params {
+ Params(std::unique_ptr<ParallelMapFunctor> parallel_map_functor,
+ int32 num_parallel_calls, bool sloppy, bool preserve_cardinality)
+ : parallel_map_functor(std::move(parallel_map_functor)),
+ num_parallel_calls(num_parallel_calls),
+ sloppy(sloppy),
+ preserve_cardinality(preserve_cardinality) {}
+
+ std::unique_ptr<ParallelMapFunctor> parallel_map_functor;
+ int32 num_parallel_calls;
+ bool sloppy;
+ bool preserve_cardinality;
+ };
+
+ ParallelMapIterator(
+ const typename DatasetBaseIterator::BaseParams& base_params,
+ const DatasetBase* input_dataset, Params params)
+ : DatasetBaseIterator(base_params),
input_dataset_(input_dataset),
- parallel_map_functor_(std::move(parallel_map_functor)),
+ parallel_map_functor_(std::move(params.parallel_map_functor)),
mu_(std::make_shared<mutex>()),
cond_var_(std::make_shared<condition_variable>()),
num_parallel_calls_(std::make_shared<model::SharedState>(
- num_parallel_calls, mu_, cond_var_)),
- sloppy_(sloppy) {
+ params.num_parallel_calls, mu_, cond_var_)),
+ sloppy_(params.sloppy),
+ preserve_cardinality_(params.preserve_cardinality) {
std::vector<string> components =
- str_util::Split(params.prefix, "::", str_util::SkipEmpty());
+ str_util::Split(base_params.prefix, "::", str_util::SkipEmpty());
prefix_end_ = components.back();
}
@@ -62,9 +76,8 @@
Status Initialize(IteratorContext* ctx) override {
mutex_lock l(*mu_);
- if (num_parallel_calls_->value == kAutoTune) {
+ if (num_parallel_calls_->value == model::kAutoTune) {
num_parallel_calls_->value = ctx->runner_threadpool_size();
- num_parallel_calls_->tunable = true;
}
TF_RETURN_IF_ERROR(
input_dataset_->MakeIterator(ctx, prefix(), &input_impl_));
@@ -86,7 +99,7 @@
RecordStop(ctx);
result->notification.WaitForNotification();
RecordStart(ctx);
- return ProcessResult(result, out_tensors, end_of_sequence);
+ return ProcessResult(ctx, result, out_tensors, end_of_sequence);
}
protected:
@@ -197,6 +210,7 @@
strings::StrCat(prefix_end_, "::active_parallel_calls"),
static_cast<float>(num_calls_));
}
+ RecordBufferEnqueue(ctx.get(), result->return_values);
result->notification.Notify();
cond_var_->notify_all();
}
@@ -225,19 +239,30 @@
&result->return_values, std::move(done));
}
- Status ProcessResult(const std::shared_ptr<InvocationResult>& result,
+ Status ProcessResult(IteratorContext* ctx,
+ const std::shared_ptr<InvocationResult>& result,
std::vector<Tensor>* out_tensors, bool* end_of_sequence)
LOCKS_EXCLUDED(*mu_) {
if (!result->end_of_input && result->status.ok()) {
*out_tensors = std::move(result->return_values);
+ RecordBufferDequeue(ctx, *out_tensors);
*end_of_sequence = false;
return Status::OK();
}
if (errors::IsOutOfRange(result->status)) {
- // `f` may deliberately raise `errors::OutOfRange` to indicate that we
- // should terminate the iteration early.
- *end_of_sequence = true;
- return Status::OK();
+ if (preserve_cardinality_) {
+ // To guarantee that the transformation preserves the cardinality of the
+ // dataset, we convert `OutOfRange` to `InvalidArgument` as the former
+ // may be interpreted by a caller as the end of sequence.
+ return errors::InvalidArgument(
+ "Function invocation produced OutOfRangeError: ",
+ result->status.error_message());
+ } else {
+ // `f` may deliberately raise `errors::OutOfRange` to indicate
+ // that we should terminate the iteration early.
+ *end_of_sequence = true;
+ return Status::OK();
+ }
}
*end_of_sequence = result->end_of_input;
return result->status;
@@ -369,6 +394,7 @@
const std::shared_ptr<model::SharedState> num_parallel_calls_;
// Determines whether outputs can be produced in non-deterministic order.
const bool sloppy_;
+ const bool preserve_cardinality_;
// Counts the number of outstanding calls.
int64 num_calls_ GUARDED_BY(*mu_) = 0;
std::unique_ptr<IteratorBase> input_impl_;
@@ -386,10 +412,12 @@
const DatasetBaseIterator::BaseParams& params,
const DatasetBase* input_dataset,
std::unique_ptr<ParallelMapFunctor> parallel_map_functor,
- int32 num_parallel_calls, bool sloppy) {
- return MakeUnique<ParallelMapIterator>(params, input_dataset,
- std::move(parallel_map_functor),
- num_parallel_calls, sloppy);
+ int32 num_parallel_calls, bool sloppy, bool preserve_cardinality) {
+ return MakeUnique<ParallelMapIterator>(
+ params, input_dataset,
+ ParallelMapIterator::Params{std::move(parallel_map_functor),
+ num_parallel_calls, sloppy,
+ preserve_cardinality});
}
} // namespace data
diff --git a/tensorflow/core/kernels/data/parallel_map_iterator.h b/tensorflow/core/kernels/data/parallel_map_iterator.h
index 08c16a6..de30446 100644
--- a/tensorflow/core/kernels/data/parallel_map_iterator.h
+++ b/tensorflow/core/kernels/data/parallel_map_iterator.h
@@ -48,7 +48,7 @@
const DatasetBaseIterator::BaseParams& params,
const DatasetBase* input_dataset,
std::unique_ptr<ParallelMapFunctor> parallel_map_functor,
- int32 num_parallel_calls, bool sloppy);
+ int32 num_parallel_calls, bool sloppy, bool preserve_cardinality);
} // namespace data
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/prefetch_dataset_op.cc b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
index 9b2a530..08d6de4 100644
--- a/tensorflow/core/kernels/data/prefetch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
@@ -125,7 +125,7 @@
}
if (!buffer_.empty()) {
- return Consume(out_tensors, end_of_sequence, ctx);
+ return Consume(ctx, out_tensors, end_of_sequence);
}
if (prefetch_thread_finished_) {
@@ -228,8 +228,8 @@
std::vector<Tensor> value;
};
- Status Consume(std::vector<Tensor>* out_tensors, bool* end_of_sequence,
- IteratorContext* ctx) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ Status Consume(IteratorContext* ctx, std::vector<Tensor>* out_tensors,
+ bool* end_of_sequence) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
const auto& stats_aggregator = ctx->stats_aggregator();
if (stats_aggregator) {
stats_aggregator->AddToHistogram(
@@ -248,6 +248,7 @@
Status s = buffer_.front().status;
if (s.ok()) {
*out_tensors = std::move(buffer_.front().value);
+ RecordBufferDequeue(ctx, *out_tensors);
}
auto_tuner_.RecordConsumption(buffer_.size());
buffer_.pop_front();
@@ -318,6 +319,7 @@
// 3. Signal that the element has been produced.
{
mutex_lock l(mu_);
+ RecordBufferEnqueue(ctx.get(), buffer_element.value);
buffer_.push_back(std::move(buffer_element));
cond_var_.notify_all();
}
diff --git a/tensorflow/core/kernels/data/shuffle_dataset_op.cc b/tensorflow/core/kernels/data/shuffle_dataset_op.cc
index 0ca347e..7134793 100644
--- a/tensorflow/core/kernels/data/shuffle_dataset_op.cc
+++ b/tensorflow/core/kernels/data/shuffle_dataset_op.cc
@@ -18,6 +18,7 @@
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/partial_tensor_shape.h"
+#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/random/philox_random.h"
#include "tensorflow/core/lib/random/random.h"
@@ -70,9 +71,9 @@
explicit Iterator(const typename DatasetIterator<T>::Params& params,
int64 seed, int64 seed2)
: DatasetIterator<T>(params),
- input_impl_(nullptr),
seed_(seed),
seed2_(seed2),
+ input_impl_(nullptr),
epoch_(0),
num_elements_(0),
parent_generator_(seed, seed2),
@@ -126,6 +127,7 @@
ctx, this->prefix(), &input_impl_));
}
if (!end_of_input_sequence) {
+ this->RecordBufferEnqueue(ctx, input_element);
buffer_[slices_.back()->end % this->dataset()->buffer_size_] =
std::move(input_element);
num_elements_++;
@@ -153,6 +155,7 @@
int64 index =
(slices_.front()->start + offset) % this->dataset()->buffer_size_;
*out_tensors = std::move(buffer_[index]);
+ this->RecordBufferDequeue(ctx, *out_tensors);
std::swap(
buffer_[index],
buffer_[slices_.front()->start % this->dataset()->buffer_size_]);
@@ -172,6 +175,14 @@
/*ratio=*/1);
}
+ void ResetRngs() EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+ // Reset the generators based on the current iterator seeds.
+ parent_generator_ = random::PhiloxRandom(seed_, seed2_);
+ generator_ = random::SingleSampleAdapter<random::PhiloxRandom>(
+ &parent_generator_);
+ generator_.Skip(num_random_samples_);
+ }
+
Status SaveInternal(IteratorStateWriter* writer) override {
mutex_lock l(mu_);
// Save state needed to restore the random number generators.
@@ -279,6 +290,10 @@
return Status::OK();
}
+ mutex mu_;
+ int64 seed_ GUARDED_BY(mu_);
+ int64 seed2_ GUARDED_BY(mu_);
+
private:
// Used to represent slices of `buffer_` that belong to different epochs.
// The invariant maintained by the implementation is: `start` <= `end`.
@@ -299,19 +314,8 @@
return out;
}
- void ResetRngs() EXCLUSIVE_LOCKS_REQUIRED(mu_) {
- // Reset the generators based on the current iterator seeds.
- parent_generator_ = random::PhiloxRandom(seed_, seed2_);
- generator_ = random::SingleSampleAdapter<random::PhiloxRandom>(
- &parent_generator_);
- generator_.Skip(num_random_samples_);
- }
-
- mutex mu_;
std::unique_ptr<std::vector<Tensor>[]> buffer_ GUARDED_BY(mu_);
std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
- int64 seed_ GUARDED_BY(mu_);
- int64 seed2_ GUARDED_BY(mu_);
int64 epoch_ GUARDED_BY(mu_);
int64 num_elements_ GUARDED_BY(mu_);
std::deque<std::unique_ptr<Slice>> slices_ GUARDED_BY(mu_);
@@ -368,7 +372,7 @@
}
private:
- // A dataset that uses a pseduorandom sequence of seeds for the iterators
+ // A dataset that uses a pseudorandom sequence of seeds for the iterators
// created from it. Used when `reshuffle_each_iteration` is true.
class ReshufflingDataset : public ShuffleDatasetBase {
public:
@@ -376,37 +380,114 @@
int64 buffer_size, int64 seed, int64 seed2, int64 count)
: ShuffleDatasetBase(ctx, input, buffer_size, count),
seed_(seed),
- seed2_(seed2),
- parent_generator_(seed, seed2),
- generator_(&parent_generator_) {}
+ seed2_(seed2) {}
string DebugString() const override {
- mutex_lock l(mu_);
return strings::StrCat("ShuffleDatasetOp(", buffer_size_, ", ", seed_,
", ", seed2_, ")::ReshufflingDataset");
}
std::unique_ptr<IteratorBase> MakeIteratorInternal(
const string& prefix) const override {
- int64 iterator_seed;
- int64 iterator_seed2;
- {
- mutex_lock l(mu_);
- iterator_seed = Random();
- iterator_seed2 = Random();
- }
- return std::unique_ptr<IteratorBase>(
- new Iterator({this, strings::StrCat(prefix, "::Shuffle")},
- iterator_seed, iterator_seed2));
+ return std::unique_ptr<IteratorBase>(new Iterator(
+ {this, strings::StrCat(prefix, "::Shuffle")}, seed_, seed2_));
}
protected:
+ class RandomSeedGenerator : public ResourceBase {
+ public:
+ RandomSeedGenerator(int64 seed, int64 seed2)
+ : seed_(seed),
+ seed2_(seed2),
+ parent_generator_(seed, seed2),
+ generator_(&parent_generator_) {}
+
+ string DebugString() override {
+ return "ReshufflingDataset::RandomSeedGenerator";
+ }
+
+ void GenerateRandomSeeds(int64* seed1, int64* seed2) {
+ mutex_lock l(mu_);
+ num_random_samples_++;
+ *seed1 = generator_();
+ num_random_samples_++;
+ *seed2 = generator_();
+ }
+
+ int64 num_random_samples() {
+ tf_shared_lock l(mu_);
+ return num_random_samples_;
+ }
+
+ void set_num_random_samples(int64 num_random_samples) {
+ mutex_lock l(mu_);
+ num_random_samples_ = num_random_samples;
+ }
+
+ void Reset() {
+ mutex_lock l(mu_);
+ // Reset the generators based on the current seeds.
+ parent_generator_ = random::PhiloxRandom(seed_, seed2_);
+ generator_ = random::SingleSampleAdapter<random::PhiloxRandom>(
+ &parent_generator_);
+ generator_.Skip(num_random_samples_);
+ }
+
+ private:
+ const int64 seed_;
+ const int64 seed2_;
+ mutex mu_;
+ random::PhiloxRandom parent_generator_ GUARDED_BY(mu_);
+ random::SingleSampleAdapter<random::PhiloxRandom> generator_
+ GUARDED_BY(mu_);
+ int64 num_random_samples_ GUARDED_BY(mu_) = 0;
+ };
+
class Iterator : public ShuffleDatasetBase::Iterator<ReshufflingDataset> {
public:
explicit Iterator(const Params& params, int64 seed, int64 seed2)
: ShuffleDatasetBase::Iterator<ReshufflingDataset>(params, seed,
seed2) {}
+ ~Iterator() override { seed_generator_->Unref(); }
+
+ Status Initialize(IteratorContext* ctx) override {
+ // Firstly, lookup or create a seed generator from the IteratorResource
+ // resource_mgr.
+ ResourceMgr* mgr = ctx->resource_mgr();
+ RandomSeedGenerator* seed_generator;
+ const string name = strings::StrCat(prefix(), "::", dataset()->name(),
+ "::RandomSeedGenerator");
+
+ int64 dataset_seed, dataset_seed2;
+ {
+ tf_shared_lock l(mu_);
+ // Ideally we'd like to hold this lock in the LookupOrCreate method,
+ // but that trips up our Deadlock detection code.
+ dataset_seed = seed_;
+ dataset_seed2 = seed2_;
+ }
+ TF_RETURN_IF_ERROR(mgr->LookupOrCreate<RandomSeedGenerator>(
+ "tf_data", name, &seed_generator,
+ [dataset_seed,
+ dataset_seed2](RandomSeedGenerator** seed_generator) {
+ // On the first iterator creation, use the original seeds from the
+ // dataset to seed a `RandomSeedGenerator` that will provide seeds
+ // for subsequent repetitions of the same dataset.
+ *seed_generator =
+ new RandomSeedGenerator(dataset_seed, dataset_seed2);
+ return Status::OK();
+ }));
+ // Now use the seed generator to update the base class Iterator seeds
+ // and random number generator with generated seeds for the current
+ // repetition.
+ mutex_lock l(mu_);
+ seed_generator->GenerateRandomSeeds(&seed_, &seed2_);
+ ResetRngs();
+ seed_generator_ = seed_generator;
+ return Status::OK();
+ }
+
protected:
std::shared_ptr<model::Node> CreateNode(
IteratorContext* ctx, model::Node::Args args) const override {
@@ -415,12 +496,10 @@
}
Status SaveInternal(IteratorStateWriter* writer) override {
- mutex_lock l(dataset()->mu_);
-
// Save RNG state of Dataset.
TF_RETURN_IF_ERROR(
writer->WriteScalar(full_name("ds_num_random_samples"),
- dataset()->num_random_samples_));
+ seed_generator_->num_random_samples()));
// Save the Iterator.
return ShuffleDatasetBase::Iterator<ReshufflingDataset>::SaveInternal(
@@ -429,24 +508,25 @@
Status RestoreInternal(IteratorContext* ctx,
IteratorStateReader* reader) override {
- mutex_lock l(dataset()->mu_);
-
// Restore RNG state of Dataset.
- TF_RETURN_IF_ERROR(
- reader->ReadScalar(full_name("ds_num_random_samples"),
- &dataset()->num_random_samples_));
- dataset()->ResetRngs();
+ int64 num_random_samples;
+ TF_RETURN_IF_ERROR(reader->ReadScalar(
+ full_name("ds_num_random_samples"), &num_random_samples));
+ seed_generator_->set_num_random_samples(num_random_samples);
+ seed_generator_->Reset();
// Restore the Iterator.
return ShuffleDatasetBase::Iterator<
ReshufflingDataset>::RestoreInternal(ctx, reader);
}
+
+ private:
+ RandomSeedGenerator* seed_generator_;
};
Status AsGraphDefInternal(SerializationContext* ctx,
DatasetGraphDefBuilder* b,
Node** output) const override {
- mutex_lock l(mu_);
Node* input_graph_node = nullptr;
TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph_node));
Node* buffer_size = nullptr;
@@ -467,28 +547,8 @@
}
private:
- random::SingleSampleAdapter<random::PhiloxRandom>::ResultType Random() const
- EXCLUSIVE_LOCKS_REQUIRED(mu_) {
- num_random_samples_++;
- auto out = generator_();
- return out;
- }
-
- void ResetRngs() const EXCLUSIVE_LOCKS_REQUIRED(mu_) {
- // Reset the generators based on the current seeds.
- parent_generator_ = random::PhiloxRandom(seed_, seed2_);
- generator_ =
- random::SingleSampleAdapter<random::PhiloxRandom>(&parent_generator_);
- generator_.Skip(num_random_samples_);
- }
-
- mutable int64 seed_ GUARDED_BY(mu_);
- mutable int64 seed2_ GUARDED_BY(mu_);
- mutable mutex mu_;
- mutable random::PhiloxRandom parent_generator_ GUARDED_BY(mu_);
- mutable random::SingleSampleAdapter<random::PhiloxRandom> generator_
- GUARDED_BY(mu_);
- mutable int64 num_random_samples_ GUARDED_BY(mu_) = 0;
+ const int64 seed_;
+ const int64 seed2_;
};
// A dataset that uses the same fixed seed for all iterators created from it.
diff --git a/tensorflow/core/kernels/data/window_dataset.cc b/tensorflow/core/kernels/data/window_dataset.cc
index bed1c66..c295631 100644
--- a/tensorflow/core/kernels/data/window_dataset.cc
+++ b/tensorflow/core/kernels/data/window_dataset.cc
@@ -41,10 +41,18 @@
return output_shapes_;
}
- string DebugString() const override { return "WindowDataset"; }
+ int64 AllocatedBytes() const override {
+ int64 allocated_bytes = 0;
+ for (auto& element : elements_) {
+ allocated_bytes += GetAllocatedBytes(element);
+ }
+ return allocated_bytes;
+ }
int64 Cardinality() const override { return elements_.size(); }
+ string DebugString() const override { return "WindowDataset"; }
+
protected:
// TODO(b/110981596): Support checkpointing.
Status AsGraphDefInternal(SerializationContext* ctx,
diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index 9d66162..ae13ae5 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -164,6 +164,7 @@
Status status =
input_impl_->GetNext(ctx, &element, end_of_sequence);
if (!*end_of_sequence) {
+ RecordBufferEnqueue(ctx, element);
buffer_.emplace_back(std::move(element), status);
} else {
input_impl_.reset();
@@ -201,8 +202,14 @@
input_impl_.reset();
}
}
+ for (size_t i = 0; i < buffer_.size(); ++i) {
+ RecordBufferDequeue(ctx, buffer_.at(i).result);
+ }
buffer_.clear();
} else {
+ for (size_t i = 0; i < window_shift; ++i) {
+ RecordBufferDequeue(ctx, buffer_.at(i).result);
+ }
buffer_.erase(buffer_.begin(), buffer_.begin() + window_shift);
}
}
diff --git a/tensorflow/core/kernels/eigen_contraction_kernel.cc b/tensorflow/core/kernels/eigen_contraction_kernel.cc
new file mode 100644
index 0000000..da42001
--- /dev/null
+++ b/tensorflow/core/kernels/eigen_contraction_kernel.cc
@@ -0,0 +1,55 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/kernels/eigen_contraction_kernel.h"
+
+#include <mutex> // NOLINT(build/c++11)
+
+// We need a pair of compile time and runtime flags to disable compilation of
+// custom contraction kernels for unsupported architectures (e.g. Android,
+// iOS, ARM and PPC CPUs, etc...), and to be able to fallback on default Eigen
+// matrix multiplication at runtime.
+//
+// It's not allowed to use absl flags library in Tensorflow, so we have to pass
+// the configuration through the environment variable.
+//
+// Example:
+// bazel test --test_env=TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL=false //test
+
+#if defined(TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL)
+
+namespace Eigen {
+namespace internal {
+
+// TODO(ezhulenev): This is a temporary workaround for disabling custom kernels
+// at runtime in tests. We should always rely on compile time flags for that.
+// Example: ... --test_env=TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL=false //test
+bool UseCustomContractionKernels() {
+ static bool use_custom_contraction_kernel = true;
+
+ static std::once_flag initialized;
+ std::call_once(initialized, [&] {
+ char* flag = std::getenv("TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL");
+ if (flag && (strcmp(flag, "false") == 0 || strcmp(flag, "0") == 0)) {
+ use_custom_contraction_kernel = false;
+ }
+ });
+
+ return use_custom_contraction_kernel;
+}
+
+} // namespace internal
+} // namespace Eigen
+#endif
diff --git a/tensorflow/core/kernels/eigen_contraction_kernel.h b/tensorflow/core/kernels/eigen_contraction_kernel.h
index 92d29e3..3d8e52c 100644
--- a/tensorflow/core/kernels/eigen_contraction_kernel.h
+++ b/tensorflow/core/kernels/eigen_contraction_kernel.h
@@ -33,11 +33,20 @@
// #endif
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
-#include "third_party/intel_mkl_dnn/include/mkldnn.h"
+
+#if defined(TENSORFLOW_USE_MKLDNN_CONTRACTION_KERNEL)
+#include "mkldnn.h"
+#endif
namespace Eigen {
namespace internal {
+#if defined(TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL)
+// Returns `true` iff we can use custom contraction kernels. This is a runtime
+// check, that uses environment variables.
+bool UseCustomContractionKernels();
+#endif // TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL
+
// Enabled by build option: "--define tensorflow_mkldnn_contraction_kernel=1"
#if defined(TENSORFLOW_USE_MKLDNN_CONTRACTION_KERNEL)
@@ -126,6 +135,11 @@
&alpha, blockA, &ldA, blockB, &ldB, &beta,
const_cast<float*>(output.data()), &ldC);
eigen_assert(st == 0);
+
+ // eigen_assert is a no-op in optimized mode so we add these to avoid
+ // compiler's unused-variable errors.
+ EIGEN_UNUSED_VARIABLE(max_index);
+ EIGEN_UNUSED_VARIABLE(st);
}
};
@@ -143,8 +157,8 @@
// Multiply default choice of block size along M and N dimensions.
// TODO(ezhulenev): Explore if this can work in general (kScaleM=2.0 worked
// well in some of models).
- static const float kScaleM = 1.5;
- static const float kScaleN = 1.0;
+ static constexpr float kScaleM = 1.5;
+ static constexpr float kScaleN = 1.0;
// Mkldnn Avx/Avx2/Avx512 unroll factors are: 8/16/48.
static const StorageIndex kUnrollM = 48;
@@ -165,6 +179,10 @@
num_threads);
}
+ // If we are using default Eigen gebp kernel there is no need to adjust the
+ // block sizes for MKL-DNN.
+ if (!UseCustomContractionKernels()) return;
+
// 2. And refine them to work well with mkldnn sgemm.
mc_ = (std::min)(
m, Eigen::divup(static_cast<StorageIndex>(mc_ * kScaleM), kUnrollM) *
@@ -206,23 +224,52 @@
typename RhsMapper::SubMapper, ColMajor>;
using GemmKernel = mkldnn_gemm_kernel<Scalar, StorageIndex, OutputMapper>;
+ // Fallback on default Eigen pack and GEBP kernel if custom contraction
+ // kernels disabled at runtime.
+ using EigenLhsPacker =
+ gemm_pack_lhs<Scalar, StorageIndex, typename LhsMapper::SubMapper,
+ Traits::mr, Traits::LhsProgress,
+ typename Traits::LhsPacket4Packing, ColMajor>;
+ using EigenRhsPacker =
+ gemm_pack_rhs<Scalar, StorageIndex, typename RhsMapper::SubMapper,
+ Traits::nr, ColMajor>;
+ using GebpKernel =
+ gebp_kernel<Scalar, Scalar, StorageIndex, OutputMapper, Traits::mr,
+ Traits::nr,
+ /*ConjugateLhs*/ false, /*ConjugateRhs*/ false>;
+
EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE static void packLhs(
Scalar* lhsBlock, const typename LhsMapper::SubMapper& data_mapper,
const StorageIndex depth, const StorageIndex rows) {
- LhsPacker()(lhsBlock, data_mapper, rows, depth);
+ if (UseCustomContractionKernels()) {
+ LhsPacker()(lhsBlock, data_mapper, rows, depth);
+ } else {
+ EigenLhsPacker()(lhsBlock, data_mapper, depth, rows, /*stride*/ 0,
+ /*offset*/ 0);
+ }
}
EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE static void packRhs(
Scalar* rhsBlock, const typename RhsMapper::SubMapper& data_mapper,
const StorageIndex depth, const StorageIndex cols) {
- RhsPacker()(rhsBlock, data_mapper, depth, cols);
+ if (UseCustomContractionKernels()) {
+ RhsPacker()(rhsBlock, data_mapper, depth, cols);
+ } else {
+ EigenRhsPacker()(rhsBlock, data_mapper, depth, cols);
+ }
}
EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE static void invoke(
const OutputMapper& output_mapper, const Scalar* lhsBlock,
const Scalar* rhsBlock, const StorageIndex rows, const StorageIndex depth,
const StorageIndex cols, const Scalar alpha) {
- GemmKernel()(output_mapper, lhsBlock, rhsBlock, rows, depth, cols, alpha);
+ if (UseCustomContractionKernels()) {
+ GemmKernel()(output_mapper, lhsBlock, rhsBlock, rows, depth, cols, alpha);
+ } else {
+ GebpKernel()(output_mapper, lhsBlock, rhsBlock, rows, depth, cols, alpha,
+ /*strideA*/ -1, /*strideB*/ -1,
+ /*offsetA*/ 0, /*offsetB*/ 0);
+ }
}
};
diff --git a/tensorflow/core/kernels/eigen_spatial_convolutions.h b/tensorflow/core/kernels/eigen_spatial_convolutions.h
index 25c735d..86d8c98 100644
--- a/tensorflow/core/kernels/eigen_spatial_convolutions.h
+++ b/tensorflow/core/kernels/eigen_spatial_convolutions.h
@@ -871,11 +871,9 @@
const bool pad_col2 = dm2.padCol(c);
const bool pad_col3 = dm3.padCol(c);
- // We can squeeze reads along the `row` and `depth` dimensions if
- // the row stride is `1`, which means that `row` and `depth`
- // dimensions are contiguous (two innermost dimensions).
- if (rhs.rowStride() == 1 && //
- !pad_col0 && !pad_col1 && !pad_col2 && !pad_col3 && //
+ // Check if we can squeeze reads along the `row` and `depth`
+ // dimensions (two innermost dimensions).
+ if (!pad_col0 && !pad_col1 && !pad_col2 && !pad_col3 && //
!dm0.padRow(start_row) && !dm0.padRow(max_row - 1) && //
!dm1.padRow(start_row) && !dm1.padRow(max_row - 1) && //
!dm2.padRow(start_row) && !dm2.padRow(max_row - 1) && //
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc
index cca3cfb..88a8a52 100644
--- a/tensorflow/core/kernels/function_ops.cc
+++ b/tensorflow/core/kernels/function_ops.cc
@@ -122,11 +122,15 @@
.TypeConstraint<string>("T"),
ArgOp);
+REGISTER_KERNEL_BUILDER(
+ Name(kArgOp).Device(DEVICE_GPU).TypeConstraint<Variant>("T"), ArgOp);
+
#define REGISTER(type) \
REGISTER_KERNEL_BUILDER( \
Name(kRetOp).Device(DEVICE_GPU).TypeConstraint<type>("T"), RetvalOp);
TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER)
TF_CALL_QUANTIZED_TYPES(REGISTER)
+REGISTER(Variant)
TF_CALL_bool(REGISTER) REGISTER_KERNEL_BUILDER(Name(kRetOp)
.Device(DEVICE_GPU)
.HostMemory("input")
@@ -134,6 +138,7 @@
RetvalOp);
REGISTER_KERNEL_BUILDER(
Name(kDeviceRetOp).Device(DEVICE_GPU).TypeConstraint<int32>("T"), RetvalOp);
+
REGISTER_KERNEL_BUILDER(Name(kRetOp)
.Device(DEVICE_GPU)
.TypeConstraint<ResourceHandle>("T")
diff --git a/tensorflow/core/kernels/fuzzing/BUILD b/tensorflow/core/kernels/fuzzing/BUILD
index 2d8b734..fcaf1a8 100644
--- a/tensorflow/core/kernels/fuzzing/BUILD
+++ b/tensorflow/core/kernels/fuzzing/BUILD
@@ -68,3 +68,5 @@
tf_oss_fuzz_corpus("decode_json_example")
tf_oss_fuzz_dict("decode_json_example")
+
+tf_ops_fuzz_target_lib("check_numerics")
diff --git a/tensorflow/core/kernels/fuzzing/check_numerics_fuzz.cc b/tensorflow/core/kernels/fuzzing/check_numerics_fuzz.cc
new file mode 100644
index 0000000..62d3989
--- /dev/null
+++ b/tensorflow/core/kernels/fuzzing/check_numerics_fuzz.cc
@@ -0,0 +1,50 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/kernels/fuzzing/fuzz_session.h"
+
+namespace tensorflow {
+namespace fuzzing {
+
+class FuzzCheckNumerics : public FuzzSession {
+ void BuildGraph(const Scope& scope) override {
+ auto input =
+ tensorflow::ops::Placeholder(scope.WithOpName("input1"), DT_FLOAT);
+ auto prefix = "Error: ";
+ (void)tensorflow::ops::CheckNumerics(scope.WithOpName("output"), input,
+ prefix);
+ }
+
+ void FuzzImpl(const uint8_t* data, size_t size) override {
+ size_t ratio = sizeof(float) / sizeof(uint8_t);
+ size_t num_floats = size / ratio;
+ const float* float_data = reinterpret_cast<const float*>(data);
+
+ Tensor input_tensor(tensorflow::DT_FLOAT,
+ TensorShape({static_cast<int64>(num_floats)}));
+ auto flat_tensor = input_tensor.flat<float>();
+ for (size_t i = 0; i < num_floats; i++) {
+ flat_tensor(i) = float_data[i];
+ }
+ RunOneInput(input_tensor).IgnoreError();
+ }
+};
+
+STANDARD_TF_FUZZ_FUNCTION(FuzzCheckNumerics);
+
+} // end namespace fuzzing
+} // end namespace tensorflow
diff --git a/tensorflow/core/kernels/lu_op.cc b/tensorflow/core/kernels/lu_op.cc
new file mode 100644
index 0000000..f9591d1
--- /dev/null
+++ b/tensorflow/core/kernels/lu_op.cc
@@ -0,0 +1,193 @@
+/* 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.
+==============================================================================*/
+
+#include "third_party/eigen3/Eigen/Core"
+#include "third_party/eigen3/Eigen/LU"
+#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/lib/math/math_util.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/work_sharder.h"
+
+namespace tensorflow {
+
+typedef Eigen::ThreadPoolDevice CPUDevice;
+
+template <typename Scalar, typename Tidx>
+class LuOp : public OpKernel {
+ public:
+ explicit LuOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+ protected:
+ using TensorShapes = gtl::InlinedVector<TensorShape, 4>;
+ using TensorOutputs = gtl::InlinedVector<Tensor*, 4>;
+
+ using Matrix =
+ Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
+ using ConstMatrixMap = Eigen::Map<const Matrix>;
+ using MatrixMap = Eigen::Map<Matrix>;
+
+ using RealScalar = typename Eigen::NumTraits<Scalar>::Real;
+
+ using Indices =
+ Eigen::Matrix<Tidx, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
+ using IndicesMap = Eigen::Map<Indices>;
+ using ConstIndicesMap = Eigen::Map<const Indices>;
+
+ public:
+ // Returns the cost per matrix operation. This is used to determine the
+ // number of threads to use for parallelizing factorization in batch mode.
+ // Cost per unit is assumed to be roughly 1ns, based on comments
+ // in core/util/work_sharder.cc.
+ // LU decomposition for a square matrix takes roughly (2/3) * (num_rows)^3.
+ // TODO(anudhyan): Refine this estimate after taking constant factors into
+ // account.
+ int64 GetCostPerUnit(const TensorShape& input_matrix_shape) const {
+ double num_rows = static_cast<double>(input_matrix_shape.dim_size(0));
+ double cost = (2 / 3.0) * MathUtil::IPow(num_rows, 3);
+ return cost >= static_cast<double>(kint64max) ? kint64max
+ : static_cast<int64>(cost);
+ }
+
+ void Compute(OpKernelContext* context) override {
+ OP_REQUIRES(context, context->num_inputs() == 1,
+ errors::InvalidArgument("Expecting exactly one input, got ",
+ context->num_inputs()));
+
+ const Tensor& input = context->input(0);
+ int input_rank = input.dims();
+ OP_REQUIRES(context, input_rank >= 2,
+ errors::InvalidArgument(
+ "Input tensor must have rank >= 2, got ", input_rank));
+
+ // If the tensor rank is greater than 2, we consider the inner-most
+ // dimensions as matrices, and loop over all the other outer ("batch")
+ // dimensions to compute the results.
+ TensorShape input_matrix_shape;
+ TensorShape batch_shape;
+ for (int dim = 0; dim < input_rank - 2; ++dim) {
+ batch_shape.AddDim(input.dim_size(dim));
+ }
+ const int64 num_rows = input.dim_size(input_rank - 2);
+ const int64 num_cols = input.dim_size(input_rank - 1);
+
+ input_matrix_shape.AppendShape({num_rows, num_cols});
+ OP_REQUIRES(context, TensorShapeUtils::IsSquareMatrix(input_matrix_shape),
+ errors::InvalidArgument("Input matrix must be square."));
+
+ // packed_triangular_factors is a matrix with the same shape as the input;
+ // permutation is a vector.
+ TensorShape permutation_shape = batch_shape;
+ permutation_shape.AddDim(num_rows);
+
+ TensorShapes output_matrix_shapes({input.shape(), permutation_shape});
+
+ TensorOutputs outputs;
+ Tensor* output_packed_triangular_factors = nullptr;
+ OP_REQUIRES_OK(
+ context, context->forward_input_or_allocate_output(
+ {0}, 0, input.shape(), &output_packed_triangular_factors));
+ outputs.emplace_back(output_packed_triangular_factors);
+
+ Tensor* output_permutation = nullptr;
+ OP_REQUIRES_OK(context, context->allocate_output(1, permutation_shape,
+ &output_permutation));
+ outputs.emplace_back(output_permutation);
+
+ if (num_rows == 0) {
+ return;
+ }
+
+ // Process the individual matrix problems in parallel using a threadpool.
+ auto shard = [this, &input, &num_rows, &num_cols, &outputs,
+ &output_matrix_shapes, context](int64 begin, int64 end) {
+ for (int64 i = begin; i < end; ++i) {
+ ComputeTensorSlice(context, i, input, num_rows, num_cols, outputs,
+ output_matrix_shapes);
+ }
+ };
+ auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads());
+ Shard(worker_threads.num_threads, worker_threads.workers,
+ batch_shape.num_elements(), GetCostPerUnit(input_matrix_shape),
+ shard);
+ }
+
+ void ComputeTensorSlice(OpKernelContext* context, int64 matrix_index,
+ const Tensor& input, int64 num_rows, int64 num_cols,
+ const TensorOutputs& outputs,
+ const TensorShapes& output_matrix_shapes) {
+ // TODO(kalakris): Handle alignment if possible. Eigen::Map is
+ // unaligned by default.
+ ConstMatrixMap input_matrix(
+ input.flat<Scalar>().data() + matrix_index * num_rows * num_cols,
+ num_rows, num_cols);
+
+ // packed_triangular_factors has shape [num_rows, num_cols]
+ MatrixMap packed_triangular_factors(
+ outputs[0]->flat<Scalar>().data() + matrix_index * num_rows * num_cols,
+ num_rows, num_rows);
+
+ // permutation has shape [num_rows, 1]
+ IndicesMap permutation_indices(
+ outputs[1]->flat<Tidx>().data() + matrix_index * num_rows, num_rows, 1);
+
+ Eigen::PartialPivLU<Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>
+ lu_decomposition(input_matrix);
+
+ // Output the packed triangular factors in a dense form.
+ // The lower triangular factor L corresponds to the strictly lower
+ // triangular part of packed_triangular_factors with an implicit unit
+ // diagonal. The upper triangular factor U is the upper triangular part of
+ // packed_triangular_factors. The triangular factors satisfy the equation
+ // P * input_matrix = L * U
+ // where P is the permutation matrix corresponding to the indices in
+ // permutation_indices.
+ packed_triangular_factors = lu_decomposition.matrixLU();
+ // Output the permutation matrix used for pivoting.
+ Eigen::PermutationMatrix<-1, -1, Tidx> permutation =
+ lu_decomposition.permutationP().transpose();
+ permutation_indices = permutation.indices();
+
+ // PartialPivLU cannot give strong guarantees on invertibility,
+ // but we can at least guard against exact zero pivots. This can occur as
+ // a result of basic user mistakes such providing integer valued
+ // matrices that are exactly singular, or due to underflow if this
+ // code is run with denormals being flushed to zero.
+ const RealScalar min_abs_pivot =
+ packed_triangular_factors.diagonal().cwiseAbs().minCoeff();
+ OP_REQUIRES(context, min_abs_pivot > RealScalar(0),
+ errors::InvalidArgument("Input is not invertible."));
+ }
+};
+
+#define REGISTER_LU(type, idx_type) \
+ REGISTER_KERNEL_BUILDER(Name("Lu") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .TypeConstraint<idx_type>("output_idx_type"), \
+ LuOp<type, idx_type>);
+
+REGISTER_LU(float, int32);
+REGISTER_LU(double, int32);
+REGISTER_LU(complex64, int32);
+REGISTER_LU(complex128, int32);
+
+REGISTER_LU(float, int64);
+REGISTER_LU(double, int64);
+REGISTER_LU(complex64, int64);
+REGISTER_LU(complex128, int64);
+
+} // namespace tensorflow
diff --git a/tensorflow/core/kernels/lu_op_gpu.cu.cc b/tensorflow/core/kernels/lu_op_gpu.cu.cc
new file mode 100644
index 0000000..f83744b
--- /dev/null
+++ b/tensorflow/core/kernels/lu_op_gpu.cu.cc
@@ -0,0 +1,275 @@
+/* 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+
+#include <algorithm>
+#include <vector>
+
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/kernels/cuda_solvers.h"
+#include "tensorflow/core/kernels/transpose_functor.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/cuda_kernel_helper.h"
+
+namespace tensorflow {
+
+typedef Eigen::GpuDevice GPUDevice;
+
+namespace {
+template <typename Scalar>
+__device__ void ComputePermutationFromTranspositions(
+ int64 num_rows, const int* pivots, Scalar* permutation_indices) {
+ // Fill in the output array with the identity permutation.
+ for (int i = 0; i < num_rows; ++i) {
+ permutation_indices[i] = Scalar(i);
+ }
+
+ // Compute the permutation from a sequence of transpositions encoded
+ // in the pivot array by applying the transpositions in order on the
+ // identity permutation.
+ for (int i = 0; i < num_rows; ++i) {
+ // Note: Internally, the cuBlas code uses Fortran convention (1-based)
+ // indexing so ith row was swapped with (pivots[i]-1)'th row in 0-based
+ // indexing.
+ Scalar t = permutation_indices[i];
+ permutation_indices[i] = permutation_indices[pivots[i] - 1];
+ permutation_indices[pivots[i] - 1] = t;
+ }
+}
+} // namespace
+
+// Kernel to compute the inverse of a permutation from a sequence of
+// transpositions.
+template <typename Scalar>
+__global__ void ComputePermutationFromTranspositionsKernel(
+ CudaLaunchConfig config, const int64 num_rows, const int* all_pivots,
+ Scalar* all_permutation_indices) {
+ // We only parallelize over batches here. Performance is not critical,
+ // since this cheap O(num_rows) kernel always follows an O(num_rows^3)
+ // LU factorization.
+ CUDA_1D_KERNEL_LOOP(index, config.virtual_thread_count) {
+ ComputePermutationFromTranspositions(
+ num_rows, all_pivots + index * num_rows,
+ all_permutation_indices + index * num_rows);
+ }
+}
+
+template <class Scalar, class Tidx>
+class LuOpGpu : public AsyncOpKernel {
+ public:
+ explicit LuOpGpu(OpKernelConstruction* context) : AsyncOpKernel(context) {}
+
+ void ComputeAsync(OpKernelContext* context, DoneCallback done) final {
+ const Tensor& input = context->input(0);
+
+ // Analyze shape and validate inputs.
+ const int input_rank = input.dims();
+
+ OP_REQUIRES_ASYNC(
+ context, input_rank >= 2,
+ errors::InvalidArgument("Input must have rank >= 2, got ", input_rank),
+ done);
+
+ const int64 num_rows = input.dim_size(input_rank - 2);
+ const int64 num_cols = input.dim_size(input_rank - 1);
+
+ OP_REQUIRES_ASYNC(
+ context, num_rows == num_cols,
+ errors::InvalidArgument("Input matrices must be squares, got", num_rows,
+ " != ", num_cols),
+ done);
+
+ TensorShape batch_shape;
+ for (int dim = 0; dim < input_rank - 2; ++dim) {
+ batch_shape.AddDim(input.dim_size(dim));
+ }
+ TensorShape permutation_indices_shape = batch_shape;
+ permutation_indices_shape.AddDim(num_rows);
+
+ const GPUDevice& device = context->eigen_device<GPUDevice>();
+ auto solver = absl::make_unique<CudaSolver>(context);
+
+ // We output the packed triangular factors in a dense form.
+ // The lower triangular factor L corresponds to the strictly lower
+ // triangular part of packed_triangular_factors with an implicit unit
+ // diagonal. The upper triangular factor U is the upper triangular part of
+ // packed_triangular_factors. The triangular factors satisfy the equation
+ // P * input_matrix = L * U
+ // where P is the permutation matrix corresponding to the indices in
+ // permutation_indices.
+ //
+ // Reuse the input buffer or make a copy for the factorization step,
+ // depending on whether this ops owns it exclusively.
+ Tensor* packed_triangular_factors;
+ OP_REQUIRES_OK_ASYNC(context,
+ context->forward_input_or_allocate_output(
+ {0}, 0, input.shape(), &packed_triangular_factors),
+ done);
+ if (!packed_triangular_factors->SharesBufferWith(input)) {
+ device.memcpy(packed_triangular_factors->flat<Scalar>().data(),
+ input.flat<Scalar>().data(),
+ input.NumElements() * sizeof(Scalar));
+ }
+
+ // Allocate output permutation.
+ Tensor* permutation_indices = nullptr;
+ OP_REQUIRES_OK_ASYNC(context,
+ context->allocate_output(1, permutation_indices_shape,
+ &permutation_indices),
+ done);
+
+ if (input.NumElements() == 0) {
+ done();
+ return;
+ }
+
+ // Allocate a temporary Tensor to store the transposed packed triangular
+ // factors.
+ Tensor packed_triangular_factors_transpose;
+ OP_REQUIRES_OK_ASYNC(
+ context,
+ context->allocate_temp(DataTypeToEnum<Scalar>::value, input.shape(),
+ &packed_triangular_factors_transpose),
+ done);
+ auto packed_triangular_factors_transpose_reshaped =
+ packed_triangular_factors_transpose
+ .template flat_inner_dims<Scalar, 3>();
+ const int64 batch_size =
+ packed_triangular_factors_transpose_reshaped.dimension(0);
+
+ // Allocate pivots on the device.
+ Tensor pivots;
+ OP_REQUIRES_OK_ASYNC(context,
+ solver->allocate_scoped_tensor(
+ DataTypeToEnum<int32>::value,
+ TensorShape{batch_size, num_rows}, &pivots),
+ done);
+ auto pivots_mat = pivots.template matrix<int32>();
+
+ // Transpose the input. This is necessary because cuBLAS assumes
+ // column-major storage while TensorFlow uses row-major.
+ OP_REQUIRES_OK_ASYNC(
+ context,
+ DoMatrixTranspose(device, *packed_triangular_factors,
+ &packed_triangular_factors_transpose),
+ done);
+
+ std::vector<DeviceLapackInfo> dev_info;
+ if (num_rows == num_cols && num_rows / batch_size <= 128) {
+ // For small matrices or large batch sizes, we use the batched
+ // interface from cuBlas.
+ auto packed_triangular_factors_ptrs = solver->GetScratchSpace<uint8>(
+ sizeof(Scalar*) * batch_size, "packed_triangular_factors_ptrs",
+ /* on_host */ true);
+ const Scalar** packed_triangular_factors_ptrs_base =
+ reinterpret_cast<const Scalar**>(
+ packed_triangular_factors_ptrs.mutable_data());
+ for (int batch = 0; batch < batch_size; ++batch) {
+ packed_triangular_factors_ptrs_base[batch] =
+ &packed_triangular_factors_transpose_reshaped(batch, 0, 0);
+ }
+ dev_info.push_back(
+ solver->GetDeviceLapackInfo(batch_size, "getrfBatched"));
+ OP_REQUIRES_OK_ASYNC(
+ context,
+ solver->GetrfBatched(num_rows, packed_triangular_factors_ptrs_base,
+ num_rows, pivots_mat.data(), &dev_info.back(),
+ batch_size),
+ done);
+ } else {
+ // For small batch sizes we use the non-batched interface from cuSolver,
+ // which is much faster for large matrices.
+ dev_info.push_back(solver->GetDeviceLapackInfo(batch_size, "getrf"));
+ for (int batch = 0; batch < batch_size; ++batch) {
+ OP_REQUIRES_OK_ASYNC(
+ context,
+ solver->Getrf(
+ num_rows, num_cols,
+ &packed_triangular_factors_transpose_reshaped(batch, 0, 0),
+ num_rows, &pivots_mat(batch, 0), &dev_info.back()(batch)),
+ done);
+ }
+ }
+
+ // Transpose the result since we had transposed the input.
+ OP_REQUIRES_OK_ASYNC(
+ context,
+ DoMatrixTranspose(device, packed_triangular_factors_transpose,
+ packed_triangular_factors),
+ done);
+
+ // Pivots encode the permutation of the rows as a sequences of row swaps.
+ // For each index i, row i is swapped with row pivots[i].
+ int* pivots_ptr = pivots.flat<int>().data();
+ Tidx* permutation_indices_ptr =
+ permutation_indices->template flat<Tidx>().data();
+ CudaLaunchConfig cfgPivots = GetCudaLaunchConfig(batch_size, device);
+ ComputePermutationFromTranspositionsKernel<<<cfgPivots.block_count,
+ cfgPivots.thread_per_block, 0,
+ device.stream()>>>(
+ cfgPivots, num_rows, pivots_ptr, permutation_indices_ptr);
+
+ // Callback for checking info after kernels finish. Also capture the
+ // temporary Tensors/ScratchSpace so they don't get deallocated before the
+ // kernels run.
+ // TODO(rmlarsen): Use move capture once C++14 becomes available.
+ auto info_checker = [context, done, dev_info](
+ const Status& status,
+ const std::vector<HostLapackInfo>& host_infos) {
+ if (!status.ok() && errors::IsInvalidArgument(status) &&
+ !host_infos.empty()) {
+ for (int i = 0; i < host_infos[0].size(); ++i) {
+ // Match the CPU error message for singular matrices. Otherwise
+ // just print the original error message from the status below.
+ OP_REQUIRES_ASYNC(context, host_infos[0].data()[i] <= 0,
+ errors::InvalidArgument("Input is not invertible."),
+ done);
+ }
+ }
+ OP_REQUIRES_OK_ASYNC(context, status, done);
+ done();
+ };
+
+ CudaSolver::CheckLapackInfoAndDeleteSolverAsync(std::move(solver), dev_info,
+ std::move(info_checker));
+ }
+};
+
+#define REGISTER_LU_GPU(type, idx_type) \
+ REGISTER_KERNEL_BUILDER(Name("Lu") \
+ .Device(DEVICE_GPU) \
+ .TypeConstraint<type>("T") \
+ .TypeConstraint<idx_type>("output_idx_type"), \
+ LuOpGpu<type, idx_type>);
+
+REGISTER_LU_GPU(float, int32);
+REGISTER_LU_GPU(double, int32);
+REGISTER_LU_GPU(complex64, int32);
+REGISTER_LU_GPU(complex128, int32);
+
+REGISTER_LU_GPU(float, int64);
+REGISTER_LU_GPU(double, int64);
+REGISTER_LU_GPU(complex64, int64);
+REGISTER_LU_GPU(complex128, int64);
+} // namespace tensorflow
+
+#endif // GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/mkl_conv_ops.cc b/tensorflow/core/kernels/mkl_conv_ops.cc
index 4b0ced3..029c539 100644
--- a/tensorflow/core/kernels/mkl_conv_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_ops.cc
@@ -850,7 +850,8 @@
// Base class for convolution forward operations
template <typename Device, typename Tinput, typename Tfilter, typename Tbias,
- typename Toutput, typename Ttemp_output, bool biasEnabled>
+ typename Toutput, typename Ttemp_output, typename Tpadding,
+ bool biasEnabled, bool padEnabled>
class MklConvOp : public OpKernel {
public:
~MklConvOp() {}
@@ -928,6 +929,11 @@
dilations, strides;
memory::dims dst_dims_tf_order, dst_dims_mkl_order;
+ // If pad with conv2d fusion is enabled
+ if (padEnabled) {
+ PadWithConvFusion(context, padding_left, padding_right);
+ }
+
// Get shapes of input tensors in MKL-DNN order
MklDnnConvUtil conv_utl(context, strides_, padding_, data_format_,
dilations_);
@@ -936,7 +942,7 @@
conv_utl.GetConvFwdSizesInMklOrder(
src_tf_shape, filter_tf_shape, &src_dims, &filter_dims, &strides,
&dilations, &dst_dims_tf_order, &dst_dims_mkl_order, &padding_left,
- &padding_right);
+ &padding_right, padEnabled);
if (!context->status().ok()) return;
// Check for corner case - if there is nothing to compute, return.
@@ -967,7 +973,12 @@
}
bool isConv2D = (strides_.size() == 4);
-
+ // TODO(Intel-tf) Add check to make sure padEnabled is true only for 2D
+ if (!isConv2D) {
+ OP_REQUIRES(
+ context, !padEnabled,
+ errors::InvalidArgument("Pad+Conv fusion only works for 2D"));
+ }
// Create memory for user data.
// Describe how the inputs and outputs of Convolution look like. Also
// specify buffers containing actual input and output data.
@@ -1011,7 +1022,7 @@
// get a conv2d fwd from primitive pool
MklConvFwdPrimitive<float, Tinput, Tfilter, Tbias, Ttemp_output>*
conv_fwd = nullptr;
- if (biasEnabled) {
+ if (fuse_biasadd_) {
memory::dims bias_dims = {};
conv_utl.GetBiasSizeInMklOrder(kInputIndex_Bias, &bias_dims);
MklConvFwdParams convFwdDims(src_dims, filter_dims, bias_dims,
@@ -1083,7 +1094,7 @@
}
// execute convolution
- if (biasEnabled) {
+ if (fuse_biasadd_) {
const Tensor& bias_tensor = MklGetInput(context, kInputIndex_Bias);
Tbias* bias_data =
this->GetBiasHandle(context, conv_fwd_pd, bias_tensor);
@@ -1104,7 +1115,51 @@
}
}
+ void PadWithConvFusion(OpKernelContext* context, memory::dims& padding_left,
+ memory::dims& padding_right) {
+ const Tensor& paddings_tf = MklGetInput(context, 2);
+ OP_REQUIRES(context, paddings_tf.dims() == 2,
+ errors::InvalidArgument("paddings must be 2-dimensional: ",
+ paddings_tf.shape().DebugString()));
+ Tpadding* paddings = nullptr;
+ // To get individual pad, need to flatten the tensor
+ paddings = static_cast<Tpadding*>(
+ const_cast<Tpadding*>(paddings_tf.flat<Tpadding>().data()));
+ // For NHWC format:
+ // paddings[0], paddings[1], paddings[6], paddings[7] should be zero
+ // if the paddings_tf is [ [0, 0] [1,2] [3,4] [0,0] ]
+ // paddings = {0, 0, 1, 2, 3, 4, 0, 0} ; flat method is row major
+ // then, values are: top = 1, bottom =2, left=3, right=4
+ // For NCHW format:
+ // paddings[0], paddings[1], paddings[2], paddings[3] should be zero
+ // similar explanation as NHWC format will apply.
+ int64 pad_top, pad_left;
+ int64 pad_bottom, pad_right;
+ string data_format = ToString(data_format_);
+ if (data_format == "NHWC") {
+ pad_top = paddings[2];
+ pad_bottom = paddings[3];
+ pad_left = paddings[4];
+ pad_right = paddings[5];
+ } else if (data_format == "NCHW") {
+ pad_top = paddings[4];
+ pad_bottom = paddings[5];
+ pad_left = paddings[6];
+ pad_right = paddings[7];
+ }
+ // Create padding arrays for MKL DNN convolutions.
+ // MKL-DNN uses asymetric padding.
+ padding_left = {static_cast<int>(pad_top), static_cast<int>(pad_left)};
+ padding_right = {static_cast<int>(pad_bottom), static_cast<int>(pad_right)};
+ }
+
protected:
+ void set_fuse_biasadd(bool fuse_biasadd) { fuse_biasadd_ = fuse_biasadd; }
+ void set_fuse_relu(bool fuse_relu) { fuse_relu_ = fuse_relu; }
+
+ // This method is for the base class MklConvOp, which handles the
+ // floating point implementation of Conv. The quantized conv implementations
+ // will use overidden versions of this method.
virtual void ExtendConvFwdParams(OpKernelContext* context,
MklConvFwdParams& params) {
// Create a string from data types of input, filter, bias, and output.
@@ -1112,6 +1167,11 @@
params.dtypes.append(typeid(Tfilter).name());
params.dtypes.append(typeid(Tbias).name());
params.dtypes.append(typeid(Toutput).name());
+
+ // Add fusions as post ops
+ // Note: Fusion of BiasAdd is handled directly inside MklConvOp by
+ // checking fuse_biasadd_ flag.
+ if (fuse_relu_) params.post_op_params.push_back({"relu", {1.0, 0.0, 0.0}});
}
virtual Tbias* GetBiasHandle(
@@ -1119,7 +1179,7 @@
std::shared_ptr<mkldnn::convolution_forward::primitive_desc>&
conv2d_fwd_pd,
const Tensor& bias_tensor) {
- if (biasEnabled) {
+ if (fuse_biasadd_) {
return static_cast<Tbias*>(
const_cast<Tbias*>(bias_tensor.flat<Tbias>().data()));
} else {
@@ -1165,7 +1225,13 @@
std::vector<int32> dilations_;
Padding padding_;
TensorFormat data_format_;
+
+ // Initialize to values the template is instantiated with
+ bool fuse_biasadd_ = biasEnabled;
+ bool fuse_relu_ = false;
+
const int kInputIndex_Src = 0, kInputIndex_Filter = 1, kInputIndex_Bias = 2;
+ const int kInputIndex_Pad = 2;
const int kOutputIndex_Dst = 0, kOutputIndex_Filter = 1;
const int kDilationH = 0, kDilationW = 1;
@@ -1217,12 +1283,12 @@
// Create convolution primitive and add it to net.
std::vector<primitive> net;
if (bias) {
- DCHECK(biasEnabled);
+ DCHECK(fuse_biasadd_);
net.push_back(convolution_forward(conv_prim_desc, src->GetOpMem(),
filter->GetOpMem(), bias->GetOpMem(),
output->GetOpMem()));
} else {
- DCHECK(!biasEnabled);
+ DCHECK(!fuse_biasadd_);
net.push_back(convolution_forward(conv_prim_desc, src->GetOpMem(),
filter->GetOpMem(),
output->GetOpMem()));
@@ -1232,13 +1298,56 @@
}
};
+// Base class for fused convolution forward operations
+template <typename Device, typename Tinput, typename Tfilter, typename Tbias,
+ typename Toutput, typename Ttemp_output>
+class MklFusedConvOp : public MklConvOp<Device, Tinput, Tfilter, Tbias, Toutput,
+ Ttemp_output, int32, false, false> {
+ public:
+ explicit MklFusedConvOp(OpKernelConstruction* context)
+ : MklConvOp<Device, Tinput, Tfilter, Tbias, Toutput, Ttemp_output, int32,
+ false, false>(context) {
+ // Since we came here through the registration of _MklFusedConv2D, get
+ // all information from 'fused_ops' and 'num_args'
+ std::vector<string> fused_ops;
+ OP_REQUIRES_OK(context, context->GetAttr("fused_ops", &fused_ops));
+
+ int num_args;
+ OP_REQUIRES_OK(context, context->GetAttr("num_args", &num_args));
+ OP_REQUIRES(context, !fused_ops.empty(),
+ errors::InvalidArgument(
+ "Fused Conv2D must have at least one fused op."));
+
+ if (fused_ops == std::vector<string>{"BiasAdd"}) {
+ this->set_fuse_biasadd(true);
+ OP_REQUIRES(context, num_args == 1,
+ errors::InvalidArgument(
+ "Fused Conv2D must have one extra argument: bias."));
+ } else if (fused_ops == std::vector<string>{"Relu"}) {
+ this->set_fuse_relu(true);
+ } else if (fused_ops == std::vector<string>{"BiasAdd", "Relu"}) {
+ this->set_fuse_biasadd(true);
+ this->set_fuse_relu(true);
+ OP_REQUIRES(context, num_args == 1,
+ errors::InvalidArgument(
+ "Fused Conv2D must have one extra argument: bias."));
+ } else {
+ OP_REQUIRES(context, false,
+ errors::Unimplemented("Fusion is not implemented: [",
+ str_util::Join(fused_ops, ","), "]"));
+ }
+ }
+
+ virtual ~MklFusedConvOp() {}
+};
+
// We create new class for each verison of Quantized Convolution and inherit
// from the FP32 version of the base class
template <typename Device, typename Tbias, typename Toutput,
typename Ttemp_output, bool biasEnabled>
class MklQuantizedConv2DOp
: public MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output,
- biasEnabled> {
+ int32, biasEnabled, false> {
public:
virtual ~MklQuantizedConv2DOp() {
if (this->input_bias_ != nullptr) {
@@ -1253,13 +1362,13 @@
}
explicit MklQuantizedConv2DOp(OpKernelConstruction* context)
- : MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output,
- biasEnabled>(context) {}
+ : MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output, int32,
+ biasEnabled, false>(context) {}
void Compute(OpKernelContext* context) override {
// Compute int32 output tensor
- MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output,
- biasEnabled>::Compute(context);
+ MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output, int32,
+ biasEnabled, false>::Compute(context);
// Compute additional outputs: min/max scalars.
int bias_index_offset;
@@ -1305,8 +1414,8 @@
protected:
void ExtendConvFwdParams(OpKernelContext* context,
MklConvFwdParams& params) override {
- MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output,
- biasEnabled>::ExtendConvFwdParams(context, params);
+ MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output, int32,
+ biasEnabled, false>::ExtendConvFwdParams(context, params);
// When the output type is quint8, the output data id requantized
// into quint8. A post_op "output_scale" is added to do the conversion.
@@ -1517,11 +1626,11 @@
}
}
// TODO(mdfaijul): Add cleaner code for non-mkl tensor
- MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output,
- biasEnabled>::AllocateOutputTensor(context, conv_prim_desc,
- output_dims_mkl_order,
- output_tf_format,
- output_tensor);
+ MklConvOp<Device, quint8, qint8, Tbias, Toutput, Ttemp_output, int32,
+ biasEnabled, false>::AllocateOutputTensor(context, conv_prim_desc,
+ output_dims_mkl_order,
+ output_tf_format,
+ output_tensor);
const Tensor& summand = MklGetInput(context, summand_idx);
if (summand.dtype() != DT_FLOAT)
TF_CHECK_OK(Status(error::Code::FAILED_PRECONDITION,
@@ -1790,34 +1899,65 @@
#endif // INTEL_MKL_ML
// Register 2D operations
-#define REGISTER_MKL_CPU_2D(T) \
- REGISTER_KERNEL_BUILDER( \
- Name("_MklConv2D") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<float>("T") \
- .Label(mkl_op_registry::kMklOpLabel), \
- MklConvOp<CPUDevice, float, float, float, float, float, false>); \
- REGISTER_KERNEL_BUILDER( \
- Name("_MklConv2DWithBias") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<float>("T") \
- .Label(mkl_op_registry::kMklOpLabel), \
- MklConvOp<CPUDevice, float, float, float, float, float, true>); \
- REGISTER_KERNEL_BUILDER(Name("__MklDummyConv2DWithBias") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<T>("T") \
- .Label(mkl_op_registry::kMklOpLabel), \
+#define REGISTER_MKL_CPU_2D(T) \
+ REGISTER_KERNEL_BUILDER(Name("_MklConv2D") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklConvOp<CPUDevice, float, float, float, float, \
+ float, int32, false, false>); \
+ REGISTER_KERNEL_BUILDER(Name("_MklConv2DWithBias") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklConvOp<CPUDevice, float, float, float, float, \
+ float, int32, true, false>); \
+ REGISTER_KERNEL_BUILDER(Name("__MklDummyConv2DWithBias") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklDummyOp<CPUDevice, T>); \
+ REGISTER_KERNEL_BUILDER(Name("_MklPadWithConv2D") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .TypeConstraint<int32>("Tpaddings") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklConvOp<CPUDevice, float, float, float, float, \
+ float, int32, false, true>); \
+ REGISTER_KERNEL_BUILDER(Name("_MklPadWithConv2D") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .TypeConstraint<int64>("Tpaddings") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklConvOp<CPUDevice, float, float, float, float, \
+ float, int64, false, true>); \
+ REGISTER_KERNEL_BUILDER(Name("__MklDummyPadWithConv2D") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .TypeConstraint<int32>("Tpaddings") \
+ .Label(mkl_op_registry::kMklOpLabel), \
MklDummyOp<CPUDevice, T>);
TF_CALL_float(REGISTER_MKL_CPU_2D);
-// Register 3D operations
-#define REGISTER_MKL_CPU_3D(T) \
- REGISTER_KERNEL_BUILDER(Name("_MklConv3D") \
+#define REGISTER_MKL_CPU_2D_FUSED(T) \
+ REGISTER_KERNEL_BUILDER(Name("_MklFusedConv2D") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.Label(mkl_op_registry::kMklOpLabel), \
- MklConvOp<CPUDevice, T, T, T, T, T, false>);
+ MklFusedConvOp<CPUDevice, T, T, T, T, T>);
+// We check the fused_ops attributes to decide if bias is enabled or not.
+
+TF_CALL_float(REGISTER_MKL_CPU_2D_FUSED);
+
+// Register 3D operations
+#define REGISTER_MKL_CPU_3D(T) \
+ REGISTER_KERNEL_BUILDER( \
+ Name("_MklConv3D") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<T>("T") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklConvOp<CPUDevice, T, T, T, T, T, int32, false, false>);
TF_CALL_float(REGISTER_MKL_CPU_3D);
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mkl_conv_ops.h b/tensorflow/core/kernels/mkl_conv_ops.h
index e6989d8..e61c20d 100644
--- a/tensorflow/core/kernels/mkl_conv_ops.h
+++ b/tensorflow/core/kernels/mkl_conv_ops.h
@@ -17,8 +17,8 @@
#define TENSORFLOW_CORE_KERNELS_MKL_CONV_OPS_H_
#include <limits>
-#include <vector>
#include <memory>
+#include <vector>
#include "mkldnn.hpp"
#include "tensorflow/core/framework/numeric_op.h"
@@ -85,7 +85,7 @@
}
// Calculate Convolution dilations
- virtual inline void GetDilationsInMklOrder(memory::dims *dilations) {
+ virtual inline void GetDilationsInMklOrder(memory::dims* dilations) {
// For now we take the dilation from the second and third dimensions only
// (we do not support dilation on the batch or depth dimension).
CHECK_NOTNULL(dilations);
@@ -288,7 +288,7 @@
const TensorShape& input_shape, const TensorShape& filter_shape,
const memory::dims& strides, const memory::dims& dilations,
memory::dims* output_dims_tf_order, memory::dims* output_dims_mkl_order,
- memory::dims* pad_l, memory::dims* pad_r) {
+ memory::dims* pad_l, memory::dims* pad_r, bool padEnabled = false) {
CHECK_NOTNULL(output_dims_tf_order);
CHECK_NOTNULL(output_dims_mkl_order);
CHECK_NOTNULL(pad_l);
@@ -373,6 +373,36 @@
padding_, &out_cols, &pad_left, &pad_right));
}
+ if (isConv2D) {
+ // Conv + pad fusion is enabled only for 2D
+ // If padEnabled, i.e., pad and conv op are fused, then
+ // all pads are already passed from pad op through
+ // *pad_l and *pad_r
+ if (padEnabled) {
+ pad_top = static_cast<int64>((*pad_l)[0]);
+ pad_left = static_cast<int64>((*pad_l)[1]);
+ pad_bottom = static_cast<int64>((*pad_r)[0]);
+ pad_right = static_cast<int64>((*pad_r)[1]);
+ // update the out_rows and out_cols based on all
+ // sides of the pads coming from pad op.
+ out_rows = out_rows + (pad_top + pad_bottom) / stride_rows;
+ out_cols = out_cols + (pad_left + pad_right) / stride_cols;
+ }
+ // Handle padding. MKL-DNN uses asymetric padding.
+ // But, if padEnabled, i.e., pad and conv op are fused,
+ // then, *pad_l and *pad_r are already set from pad op.
+ // In that case they need not set here.
+ else {
+ *pad_l = {static_cast<int>(pad_top), static_cast<int>(pad_left)};
+ *pad_r = {static_cast<int>(pad_bottom), static_cast<int>(pad_right)};
+ }
+ } else {
+ // Set padding for Conv3D here
+ *pad_l = {static_cast<int>(pad_D1), static_cast<int>(pad_top),
+ static_cast<int>(pad_left)};
+ *pad_r = {static_cast<int>(pad_D2), static_cast<int>(pad_bottom),
+ static_cast<int>(pad_right)};
+ }
// Tensorflow output is in data_format order.
// Conv2D: NHWC or NCHW
// Conv3D: NDHWC or NCDHW
@@ -393,9 +423,6 @@
mkldnn_sizes[MklDnnDims::Dim_H] = static_cast<int>(out_rows);
mkldnn_sizes[MklDnnDims::Dim_W] = static_cast<int>(out_cols);
*output_dims_mkl_order = mkldnn_sizes;
-
- *pad_l = {static_cast<int>(pad_top), static_cast<int>(pad_left)};
- *pad_r = {static_cast<int>(pad_bottom), static_cast<int>(pad_right)};
} else {
std::vector<int> mkldnn_sizes(5, -1);
mkldnn_sizes[MklDnnDims3D::Dim3d_N] = out_batch;
@@ -404,11 +431,6 @@
mkldnn_sizes[MklDnnDims3D::Dim3d_H] = static_cast<int>(out_rows);
mkldnn_sizes[MklDnnDims3D::Dim3d_W] = static_cast<int>(out_cols);
*output_dims_mkl_order = mkldnn_sizes;
-
- *pad_l = {static_cast<int>(pad_D1), static_cast<int>(pad_top),
- static_cast<int>(pad_left)};
- *pad_r = {static_cast<int>(pad_D2), static_cast<int>(pad_bottom),
- static_cast<int>(pad_right)};
}
}
@@ -441,8 +463,8 @@
input_tf_shape.DebugString()));
}
- GetOutputAndPadSizeInMklOrder(input_tf_shape, filter_tf_shape,
- strides, dilations, output_dims_tf_order,
+ GetOutputAndPadSizeInMklOrder(input_tf_shape, filter_tf_shape, strides,
+ dilations, output_dims_tf_order,
output_dims_mkl_order, pad_l, pad_r);
}
@@ -457,10 +479,9 @@
inline void GetConvFwdSizesInMklOrder(
const TensorShape& input_shape, const TensorShape& filter_shape,
memory::dims* input_dims, memory::dims* filter_dims,
- memory::dims* strides, memory::dims *dilations,
- memory::dims* output_dims_tf_order,
- memory::dims* output_dims_mkl_order, memory::dims* pad_l,
- memory::dims* pad_r) {
+ memory::dims* strides, memory::dims* dilations,
+ memory::dims* output_dims_tf_order, memory::dims* output_dims_mkl_order,
+ memory::dims* pad_l, memory::dims* pad_r, bool padEnabled = false) {
CHECK_NOTNULL(input_dims);
CHECK_NOTNULL(filter_dims);
CHECK_NOTNULL(strides);
@@ -476,10 +497,9 @@
if (!context_->status().ok()) return;
GetStridesInMklOrder(strides);
GetDilationsInMklOrder(dilations);
- GetOutputAndPadSizeInMklOrder(input_shape, filter_shape,
- *strides, *dilations,
- output_dims_tf_order, output_dims_mkl_order,
- pad_l, pad_r);
+ GetOutputAndPadSizeInMklOrder(
+ input_shape, filter_shape, *strides, *dilations, output_dims_tf_order,
+ output_dims_mkl_order, pad_l, pad_r, padEnabled);
if (!context_->status().ok()) return;
}
};
@@ -536,7 +556,6 @@
TensorFormat data_format_; // NCHW or NHWC
};
-
/////////////////////////////////////////////////////////////////////
/// Dummy Mkl op that is just used for operators that are intermediate
/// output of node fusion in the graph
diff --git a/tensorflow/core/kernels/mkl_fused_ops_test.cc b/tensorflow/core/kernels/mkl_fused_ops_test.cc
new file mode 100644
index 0000000..258cca9
--- /dev/null
+++ b/tensorflow/core/kernels/mkl_fused_ops_test.cc
@@ -0,0 +1,405 @@
+/* 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.
+==============================================================================*/
+#ifdef INTEL_MKL
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/image_ops.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/conv_ops_gpu.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/platform/types.h"
+#include "tensorflow/core/public/session.h"
+
+namespace tensorflow {
+
+// Helper class for converting MKL tensors to TF tensors and comparing to
+// expected values
+
+static const uint8 dummy_tensor[] = {0, 0, 0, 0, 0, 0, 0, 0};
+static const TensorShape dummy_shape({8});
+
+template <typename T>
+class ConvMklToTF : public OpsTestBase {
+ public:
+ void PerformConversion(DataType dtype, const Tensor& tensor,
+ const Tensor& mkl_meta_tensor, Tensor* output) {
+ // Create an MKL to TF conversion node and execute it
+ TF_EXPECT_OK(NodeDefBuilder("mkl_to_tf_op", "_MklToTf")
+ .Input(FakeInput(dtype)) // Input
+ .Input(FakeInput(DT_UINT8)) // Mkl second tensor
+ .Attr("T", dtype)
+ .Attr("_kernel", "MklOp")
+ .Finalize(node_def()));
+ TF_EXPECT_OK(InitOp());
+ AddInputFromArray<T>(tensor.shape(), tensor.flat<T>());
+ AddInputFromArray<uint8>(mkl_meta_tensor.shape(),
+ mkl_meta_tensor.flat<uint8>());
+ TF_ASSERT_OK(RunOpKernel());
+
+ *output = *GetOutput(0);
+ }
+
+ void ConvertAndCompare(DataType dtype, const Tensor& tensor,
+ const Tensor& mkl_meta_tensor,
+ const Tensor& expected) {
+ Tensor output;
+ PerformConversion(dtype, tensor, mkl_meta_tensor, &output);
+ test::ExpectTensorNear<T>(expected, output, 1e-5);
+ }
+ void TestBody() {}
+};
+
+// Testing MKL's fused convolution ops
+
+template <typename T>
+class MklFusedConv2DOpTest : public OpsTestBase {
+ protected:
+ static constexpr int kDepth = 3;
+ static constexpr int kImageWidth = 32;
+ static constexpr int kImageHeight = 32;
+ static constexpr int kImageBatchCount = 8;
+
+ using BiasAddGraphRunner =
+ std::function<void(const Tensor& input_data, const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* out)>;
+
+ // Runs a Tensorflow graph defined by the root scope, and fetches the result
+ // of 'fetch' node into the output Tensor.
+ void RunAndFetch(const tensorflow::Scope& root, const string& fetch,
+ Tensor* output) {
+ tensorflow::GraphDef graph;
+ TF_ASSERT_OK(root.ToGraphDef(&graph));
+
+ std::unique_ptr<tensorflow::Session> session(
+ tensorflow::NewSession(tensorflow::SessionOptions()));
+ TF_ASSERT_OK(session->Create(graph));
+
+ std::vector<Tensor> unfused_tensors;
+ TF_ASSERT_OK(session->Run({}, {fetch}, {}, &unfused_tensors));
+
+ *output = unfused_tensors[0];
+ }
+
+ void RunConv2DWithBias(const Tensor& input_data, const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* output,
+ int stride = 1) {
+ auto root = tensorflow::Scope::NewRootScope();
+
+ auto conv = ops::Conv2D(
+ root.WithOpName("conv"),
+ ops::Const(root.WithOpName("input"), Input::Initializer(input_data)),
+ ops::Const(root.WithOpName("filter"), Input::Initializer(filter_data)),
+ {1, stride, stride, 1}, "SAME");
+
+ auto with_bias = ops::BiasAdd(
+ root.WithOpName("with_bias"), conv,
+ ops::Const(root.WithOpName("bias"), Input::Initializer(bias_data)));
+
+ RunAndFetch(root, "with_bias", output);
+ }
+
+ void RunConv2DWithBiasAndRelu(const Tensor& input_data,
+ const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* output,
+ int stride = 1) {
+ auto root = tensorflow::Scope::NewRootScope();
+
+ auto conv = ops::Conv2D(
+ root.WithOpName("conv"),
+ ops::Const(root.WithOpName("input"), Input::Initializer(input_data)),
+ ops::Const(root.WithOpName("filter"), Input::Initializer(filter_data)),
+ {1, stride, stride, 1}, "SAME");
+
+ auto with_bias = ops::BiasAdd(
+ root.WithOpName("with_bias"), conv,
+ ops::Const(root.WithOpName("bias"), Input::Initializer(bias_data)));
+
+ auto with_relu = ops::Relu(root.WithOpName("with_relu"), with_bias);
+
+ RunAndFetch(root, "with_relu", output);
+ }
+
+ void RunMklFusedConv2DOp(const Tensor& image, const Tensor& filter,
+ const std::vector<Tensor>& args,
+ const std::vector<string>& fused_ops, Tensor* output,
+ int stride = 1) {
+ DataType dtype = DataTypeToEnum<T>::v();
+ int num_args = static_cast<int>(args.size());
+
+ TF_EXPECT_OK(NodeDefBuilder("fused_conv_op", "_MklFusedConv2D")
+ .Input(FakeInput(dtype))
+ .Input(FakeInput(dtype))
+ .Attr("num_args", num_args)
+ .Input(FakeInput(num_args, dtype))
+ .Input(FakeInput(DT_UINT8))
+ .Input(FakeInput(DT_UINT8))
+ .Input(FakeInput(num_args, DT_UINT8))
+ .Attr("T", dtype)
+ .Attr("strides", {1, stride, stride, 1})
+ .Attr("padding", "SAME")
+ .Attr("fused_ops", fused_ops)
+ .Attr("_kernel", "MklOp")
+ .Finalize(node_def()));
+
+ TF_EXPECT_OK(InitOp());
+
+ AddInputFromArray<T>(image.shape(), image.flat<T>());
+ AddInputFromArray<T>(filter.shape(), filter.flat<T>());
+ for (const Tensor& arg : args)
+ AddInputFromArray<T>(arg.shape(), arg.flat<T>());
+ AddInputFromArray<uint8>(dummy_shape, dummy_tensor);
+ AddInputFromArray<uint8>(dummy_shape, dummy_tensor);
+ for (const Tensor& arg : args)
+ AddInputFromArray<uint8>(dummy_shape, dummy_tensor);
+ TF_ASSERT_OK(RunOpKernel());
+
+ // Compare output to expected results
+ const Tensor& output_tensor = *GetOutput(0);
+ // Index 2 will need to be changed if the number of outputs produced
+ // by MklConv2D change.
+ const Tensor& output_meta_tensor = *GetOutput(2);
+ ConvMklToTF<T> conv_comp;
+ conv_comp.PerformConversion(dtype, output_tensor, output_meta_tensor,
+ output);
+ }
+
+ void VerifyBiasAddTensorsNear(int depth, int image_width, int image_height,
+ int image_batch_count, int filter_size,
+ int filter_count,
+ const BiasAddGraphRunner& run_default,
+ const BiasAddGraphRunner& run_fused) {
+ DataType dtype = DataTypeToEnum<T>::v();
+
+ Tensor image(dtype, {image_batch_count, image_height, image_width, depth});
+ image.flat<T>() = image.flat<T>().setRandom();
+
+ Tensor filter(dtype, {filter_size, filter_size, depth, filter_count});
+ filter.flat<T>() = filter.flat<T>().setRandom();
+
+ const int bias_size = filter_count;
+ Tensor bias(dtype, {bias_size});
+ bias.flat<T>() = bias.flat<T>().setRandom();
+
+ Tensor conv_2d;
+ Tensor fused_conv_2d;
+
+ run_default(image, filter, bias, &conv_2d);
+ run_fused(image, filter, bias, &fused_conv_2d);
+
+ ASSERT_EQ(conv_2d.dtype(), fused_conv_2d.dtype());
+ ASSERT_EQ(conv_2d.shape(), fused_conv_2d.shape());
+
+ test::ExpectClose(conv_2d, fused_conv_2d);
+ }
+
+ // Verifies that computing Conv2D+BiasAdd in a graph is identical to
+ // FusedConv2D.
+ void VerifyConv2DWithBias(int filter_size, int filter_count,
+ int depth = kDepth, int image_width = kImageWidth,
+ int image_height = kImageHeight,
+ int image_batch_count = kImageBatchCount) {
+ const BiasAddGraphRunner run_default =
+ [this](const Tensor& input_data, const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* out) {
+ RunConv2DWithBias(input_data, filter_data, bias_data, out);
+ };
+
+ const BiasAddGraphRunner run_fused =
+ [this](const Tensor& input_data, const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* out) {
+ RunMklFusedConv2DOp(input_data, filter_data, {bias_data}, {"BiasAdd"},
+ out);
+ };
+
+ VerifyBiasAddTensorsNear(depth, image_width, image_height,
+ image_batch_count, filter_size, filter_count,
+ run_default, run_fused);
+ }
+
+ // Verifies that computing Conv2D+BiasAdd+Relu in a graph is identical to
+ // FusedConv2D.
+ void VerifyConv2DWithBiasAndRelu(int filter_size, int filter_count,
+ int depth = kDepth,
+ int image_width = kImageWidth,
+ int image_height = kImageHeight,
+ int image_batch_count = kImageBatchCount) {
+ const BiasAddGraphRunner run_default =
+ [this](const Tensor& input_data, const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* out) {
+ RunConv2DWithBiasAndRelu(input_data, filter_data, bias_data, out);
+ };
+
+ const BiasAddGraphRunner run_fused =
+ [this](const Tensor& input_data, const Tensor& filter_data,
+ const Tensor& bias_data, Tensor* out) {
+ RunMklFusedConv2DOp(input_data, filter_data, {bias_data},
+ {"BiasAdd", "Relu"}, out);
+ };
+
+ VerifyBiasAddTensorsNear(depth, image_width, image_height,
+ image_batch_count, filter_size, filter_count,
+ run_default, run_fused);
+ }
+};
+
+template <typename T>
+class MklFusedConv2DWithBiasOpTest : public MklFusedConv2DOpTest<T> {};
+
+TYPED_TEST_CASE_P(MklFusedConv2DWithBiasOpTest);
+
+// -------------------------------------------------------------------------- //
+// Conv2D + BiasAdd + {Relu} //
+// -------------------------------------------------------------------------- //
+
+TYPED_TEST_P(MklFusedConv2DWithBiasOpTest, OneByOneConvolution) {
+ const int filter_size = 1;
+ const int filter_count = 12;
+ this->VerifyConv2DWithBias(filter_size, filter_count);
+}
+
+TYPED_TEST_P(MklFusedConv2DWithBiasOpTest, SpatialConvolution) {
+ const int filter_size = 3;
+ const int filter_count = 12;
+ this->VerifyConv2DWithBias(filter_size, filter_count);
+}
+
+TYPED_TEST_P(MklFusedConv2DWithBiasOpTest, OneByOneConvolutionAndRelu) {
+ const int filter_size = 1;
+ const int filter_count = 12;
+ this->VerifyConv2DWithBiasAndRelu(filter_size, filter_count);
+}
+
+TYPED_TEST_P(MklFusedConv2DWithBiasOpTest, SpatialConvolutionAndRelu) {
+ const int filter_size = 3;
+ const int filter_count = 12;
+ this->VerifyConv2DWithBiasAndRelu(filter_size, filter_count);
+}
+
+REGISTER_TYPED_TEST_CASE_P(MklFusedConv2DWithBiasOpTest, //
+ OneByOneConvolution, //
+ SpatialConvolution, //
+ OneByOneConvolutionAndRelu, //
+ SpatialConvolutionAndRelu);
+
+using MklFusedBiasAddDataTypes = ::testing::Types<float>;
+INSTANTIATE_TYPED_TEST_CASE_P(Test, MklFusedConv2DWithBiasOpTest,
+ MklFusedBiasAddDataTypes);
+// Testing fusion of pad and convolution
+
+class FusedPadConvOpTest : public OpsTestBase {
+ public:
+ template <typename T>
+ void Run(DataType dtype, Tensor& image, Tensor& filter, Tensor& padding,
+ Tensor& expected, const string data_format) {
+ const int stride = 1;
+
+ // Create a fused pad+conv2d node
+ TF_EXPECT_OK(NodeDefBuilder("fused_pad_conv_op", "_MklPadWithConv2D")
+ .Input(FakeInput(dtype)) // Input
+ .Input(FakeInput(dtype)) // Filter
+ .Input(FakeInput(DT_INT32)) // Padding
+ .Input(FakeInput(DT_UINT8)) // MKl second tensor
+ .Input(FakeInput(DT_UINT8)) // MKl second tensor
+ .Input(FakeInput(DT_UINT8)) // MKl second tensor
+ .Attr("padding", "VALID")
+ .Attr("data_format", data_format)
+ .Attr("T", dtype)
+ .Attr("strides", {1, stride, stride, 1})
+ .Attr("_kernel", "MklOp")
+ .Finalize(node_def()));
+ TF_EXPECT_OK(InitOp());
+
+ // Setting up inputs and execute
+ AddInputFromArray<T>(image.shape(), image.flat<T>());
+ AddInputFromArray<T>(filter.shape(), filter.flat<T>());
+ AddInputFromArray<int32>(padding.shape(), padding.flat<int32>());
+ AddInputFromArray<uint8>(dummy_shape, dummy_tensor);
+ AddInputFromArray<uint8>(dummy_shape, dummy_tensor);
+ AddInputFromArray<uint8>(dummy_shape, dummy_tensor);
+ TF_ASSERT_OK(RunOpKernel());
+
+ // Compare output to expected results
+ const Tensor& first = *GetOutput(0);
+ const Tensor& second = *GetOutput(2);
+ ConvMklToTF<T> conv_comp;
+ conv_comp.ConvertAndCompare(dtype, first, second, expected);
+ }
+};
+
+TEST_F(FusedPadConvOpTest, PaddingConvTest) {
+ const int depth = 1;
+ const int image_width = 4;
+ const int image_height = 3;
+ const int image_batch_count = 1;
+ Tensor image(DT_FLOAT, {image_batch_count, image_height, image_width, depth});
+ test::FillValues<float>(&image, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+
+ const int filter_size = 3;
+ const int filter_count = 1;
+ Tensor filter(DT_FLOAT, {filter_size, filter_size, depth, filter_count});
+ test::FillValues<float>(&filter, {1, 4, 7, 2, 5, 8, 3, 6, 9});
+
+ const int padding_height = 4;
+ const int padding_width = 2;
+ Tensor padding(DT_INT32, {padding_height, padding_width});
+ test::FillValues<int32>(&padding, {0, 0, 3, 4, 1, 2, 0, 0});
+
+ Tensor expected(DT_FLOAT, TensorShape({1, 8, 5, 1}));
+ test::FillValues<float>(
+ &expected,
+ {0, 0, 0, 0, 0, 24, 42, 60, 33, 12, 105, 150, 183, 95,
+ 32, 235, 312, 357, 178, 56, 187, 234, 261, 121, 32, 106, 126, 138,
+ 59, 12, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+
+ Run<float>(DT_FLOAT, image, filter, padding, expected, "NHWC");
+}
+
+TEST_F(FusedPadConvOpTest, PaddingConvTestNchw) {
+ const int depth = 1;
+ const int image_width = 4;
+ const int image_height = 3;
+ const int image_batch_count = 1;
+ Tensor image(DT_FLOAT, {image_batch_count, depth, image_height, image_width});
+ test::FillValues<float>(&image, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+
+ const int filter_size = 3;
+ const int filter_count = 1;
+ Tensor filter(DT_FLOAT, {filter_size, filter_size, depth, filter_count});
+ test::FillValues<float>(&filter, {1, 4, 7, 2, 5, 8, 3, 6, 9});
+
+ const int padding_height = 4;
+ const int padding_width = 2;
+ Tensor padding(DT_INT32, {padding_height, padding_width});
+ test::FillValues<int32>(&padding, {0, 0, 0, 0, 3, 4, 1, 2});
+
+ Tensor expected(DT_FLOAT, TensorShape({1, 1, 8, 5}));
+ test::FillValues<float>(
+ &expected,
+ {0, 0, 0, 0, 0, 24, 42, 60, 33, 12, 105, 150, 183, 95,
+ 32, 235, 312, 357, 178, 56, 187, 234, 261, 121, 32, 106, 126, 138,
+ 59, 12, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+
+ Run<float>(DT_FLOAT, image, filter, padding, expected, "NCHW");
+}
+} // namespace tensorflow
+#endif // INTEL_MKL
diff --git a/tensorflow/core/kernels/mkl_relu_op.cc b/tensorflow/core/kernels/mkl_relu_op.cc
index 7082136..a12fd5c 100644
--- a/tensorflow/core/kernels/mkl_relu_op.cc
+++ b/tensorflow/core/kernels/mkl_relu_op.cc
@@ -204,7 +204,7 @@
~MklEltwiseFwdPrimitiveFactory() {}
static string CreateKey(const MklEltwiseFwdParams<T>& fwdParams,
- memory::format src_fmt) {
+ memory::format src_fmt) {
string prefix = "eltwise_fwd";
FactoryKeyCreator key_creator;
key_creator.AddAsKey(prefix);
@@ -422,8 +422,8 @@
private:
static string CreateKey(const MklEltwiseBwdParams<T>& bwdParams,
- const memory::format& src_fmt,
- const memory::format& diff_dst_fmt) {
+ const memory::format& src_fmt,
+ const memory::format& diff_dst_fmt) {
string prefix = "eltwise_bwd";
FactoryKeyCreator key_creator;
key_creator.AddAsKey(prefix);
@@ -856,9 +856,9 @@
Tensor* dst_tensor = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
- {static_cast<const int>(src_index)},
- static_cast<const int>(dst_index),
- tf_shape_dst, &dst_tensor));
+ {static_cast<const int>(src_index)},
+ static_cast<const int>(dst_index),
+ tf_shape_dst, &dst_tensor));
AllocateOutputSetMklShape(context, dst_index, dnn_shape_dst);
T* dst_data = dst_tensor->flat<T>().data();
@@ -867,18 +867,19 @@
eltwise_fwd->Execute(src_data, dst_data);
} catch (mkldnn::error& e) {
string error_msg = "Status: " + std::to_string(e.status) +
- ", message: " + string(e.message) +
- ", in file " + string(__FILE__) + ":" +
- std::to_string(__LINE__);
- OP_REQUIRES_OK(context,
- errors::Aborted("Operation received an exception:",
- error_msg));
+ ", message: " + string(e.message) + ", in file " +
+ string(__FILE__) + ":" + std::to_string(__LINE__);
+ OP_REQUIRES_OK(
+ context,
+ errors::Aborted("Operation received an exception:", error_msg));
}
}
private:
engine cpu_engine = engine(engine::cpu, 0);
std::shared_ptr<relu_forward::primitive_desc> relu_fwd_pd;
+
+ protected:
float alpha_;
float beta_;
};
@@ -947,11 +948,11 @@
auto diff_dst_tf_data_format =
MklDnnDataFormatToTFDataFormat(diff_dst_mkl_data_format);
- src_dims = (src_tensor.dims() == 4)
- ? TFShapeToMklDnnDimsInNCHW(src_tensor.shape(),
- diff_dst_tf_data_format)
- : TFShapeToMklDnnDimsInNCDHW(src_tensor.shape(),
- diff_dst_tf_data_format);
+ src_dims = (src_tensor.dims() == 4)
+ ? TFShapeToMklDnnDimsInNCHW(src_tensor.shape(),
+ diff_dst_tf_data_format)
+ : TFShapeToMklDnnDimsInNCDHW(src_tensor.shape(),
+ diff_dst_tf_data_format);
src_md =
memory::desc(src_dims, MklDnnType<T>(), diff_dst_mkl_data_format);
} else {
@@ -1001,8 +1002,7 @@
// allocate diff_src tensor
MklDnnShape dnn_shape_diff_src;
TensorShape tf_shape_diff_src;
- if (dnn_shape_src.IsMklTensor() ||
- dnn_shape_diff_dst.IsMklTensor()) {
+ if (dnn_shape_src.IsMklTensor() || dnn_shape_diff_dst.IsMklTensor()) {
auto diff_src_pd = eltwise_bwd_pd->diff_src_primitive_desc();
dnn_shape_diff_src.SetMklTensor(true);
dnn_shape_diff_src.SetMklLayout(&diff_src_pd);
@@ -1012,9 +1012,10 @@
dnn_shape_src.GetSizesAsMklDnnDims(),
dnn_shape_src.GetTfDataFormat());
} else {
- dnn_shape_diff_src.SetTfLayout(dnn_shape_diff_dst.GetDimension(),
- dnn_shape_diff_dst.GetSizesAsMklDnnDims(),
- dnn_shape_diff_dst.GetTfDataFormat());
+ dnn_shape_diff_src.SetTfLayout(
+ dnn_shape_diff_dst.GetDimension(),
+ dnn_shape_diff_dst.GetSizesAsMklDnnDims(),
+ dnn_shape_diff_dst.GetTfDataFormat());
}
tf_shape_diff_src.AddDim(diff_src_pd.get_size() / sizeof(T));
} else {
@@ -1045,6 +1046,8 @@
private:
engine cpu_engine = engine(engine::cpu, 0);
std::shared_ptr<relu_forward::primitive_desc> relu_fwd_pd;
+
+ protected:
float alpha_;
float beta_;
};
@@ -1312,8 +1315,86 @@
T* out_o = diff_src_tensor->flat<T>().data();
T* user_i = const_cast<T*>(src_tensor.flat<T>().data());
T* user_g = const_cast<T*>(diff_dst_tensor.flat<T>().data());
- out_o[0] = user_g[0] * user_i[0] > 0 &&
- (user_i[0] < static_cast<T>(RELU6_UPPER_BOUND));
+ out_o[0] = user_g[0] * (user_i[0] > 0 &&
+ (user_i[0] < static_cast<T>(RELU6_UPPER_BOUND)));
+ return;
+ }
+};
+
+template <typename Device, typename T>
+class MklLeakyReluOp : public MklReluOpBase<Device, T, eltwise_relu> {
+ public:
+ ~MklLeakyReluOp() {}
+
+ explicit MklLeakyReluOp(OpKernelConstruction* context)
+ : MklReluOpBase<Device, T, eltwise_relu>(context, 0.0f, 0.0f) {
+ float alpha;
+ OP_REQUIRES_OK(context, context->GetAttr("alpha", &alpha));
+ OP_REQUIRES(
+ context, alpha <= 1,
+ errors::InvalidArgument("MKL LeakyRelu only supports alpha <= 1. "
+ "alpha is: ",
+ alpha));
+
+ this->alpha_ = alpha;
+ }
+
+ virtual void Compute_Scalar(OpKernelContext* context) {
+ const size_t src_index = 0; // index of src input tensor
+ const size_t dst_index = 0; // index of dst output tensor
+ const Tensor& src_tensor = MklGetInput(context, src_index);
+ MklDnnShape dnn_shape_src;
+ GetMklShape(context, src_index, &dnn_shape_src);
+
+ Tensor* dst_tensor = nullptr;
+ T* user_i = const_cast<T*>(src_tensor.flat<T>().data());
+ MklDnnShape dnn_shape_dst;
+ dnn_shape_dst.SetMklTensor(false);
+ AllocateOutputSetMklShape(context, dst_index, &dst_tensor,
+ src_tensor.shape(), dnn_shape_dst);
+ T* out_o = dst_tensor->flat<T>().data();
+ out_o[0] = user_i[0] >= 0 ? user_g[0] : user_g[0] * this->alpha_;
+ return;
+ }
+};
+
+template <typename Device, typename T>
+class MklLeakyReluGradOp : public MklReluGradOpBase<Device, T, eltwise_relu> {
+ public:
+ ~MklLeakyReluGradOp() {}
+
+ explicit MklLeakyReluGradOp(OpKernelConstruction* context)
+ : MklReluGradOpBase<Device, T, eltwise_relu>(context, 0.0f, 0.0f) {
+ float alpha;
+ OP_REQUIRES_OK(context, context->GetAttr("alpha", &alpha));
+ OP_REQUIRES(
+ context, alpha <= 1,
+ errors::InvalidArgument("MKL LeakyRelu only supports alpha <= 1. "
+ "alpha is: ",
+ alpha));
+
+ this->alpha_ = alpha;
+ }
+
+ 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
+ const Tensor& src_tensor = MklGetInput(context, src_index);
+ const Tensor& diff_dst_tensor = MklGetInput(context, diff_dst_index);
+ Tensor* diff_src_tensor = nullptr;
+
+ MklDnnShape dnn_shape_diff_dst;
+ GetMklShape(context, diff_dst_index, &dnn_shape_diff_dst);
+
+ MklDnnShape dnn_shape_diff_src;
+ dnn_shape_diff_src.SetMklTensor(false);
+ AllocateOutputSetMklShape(context, diff_src_index, &diff_src_tensor,
+ diff_dst_tensor.shape(), dnn_shape_diff_src);
+ T* out_o = diff_src_tensor->flat<T>().data();
+ T* user_i = const_cast<T*>(src_tensor.flat<T>().data());
+ T* user_g = const_cast<T*>(diff_dst_tensor.flat<T>().data());
+ out_o[0] = user_i[0] >= 0 ? user_g[0] : user_g[0] * this->alpha_;
return;
}
};
@@ -1376,6 +1457,19 @@
MklRelu6GradOp<CPUDevice, type>);
TF_CALL_float(REGISTER_RELU6_MKL_SUPPORTED_KERNELS_TYPES);
+#define REGISTER_LeakyRelu_MKL_SUPPORTED_KERNELS_TYPES(type) \
+ REGISTER_KERNEL_BUILDER(Name("_MklLeakyRelu") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklLeakyReluOp<CPUDevice, type>); \
+ REGISTER_KERNEL_BUILDER(Name("_MklLeakyReluGrad") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .Label(mkl_op_registry::kMklOpLabel), \
+ MklLeakyReluGradOp<CPUDevice, type>);
+TF_CALL_float(REGISTER_LeakyRelu_MKL_SUPPORTED_KERNELS_TYPES);
+
#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mkl_softmax_op.cc b/tensorflow/core/kernels/mkl_softmax_op.cc
index 094129a..dc3ae3d 100644
--- a/tensorflow/core/kernels/mkl_softmax_op.cc
+++ b/tensorflow/core/kernels/mkl_softmax_op.cc
@@ -50,8 +50,6 @@
// src_tensor now points to the 0-th input of global data struct "context"
size_t src_idx = 0;
const Tensor& src_tensor = MklGetInput(context, src_idx);
- const int input_dims = src_tensor.dims();
-
// Add: get MklShape
MklDnnShape src_mkl_shape;
GetMklShape(context, src_idx, &src_mkl_shape);
@@ -61,15 +59,27 @@
auto src_tf_shape = src_mkl_shape.IsMklTensor()
? src_mkl_shape.GetTfShape()
: src_tensor.shape();
+ const int input_dims = src_tf_shape.dims();
auto src_dims = TFShapeToMklDnnDims(src_tf_shape);
- auto output_dims = src_dims;
+ memory::dims output_dims;
+ int axis;
+ if (src_mkl_shape.IsMklTensor()) {
+ axis = 1;
+ output_dims = src_mkl_shape.GetSizesAsMklDnnDims();
+ } else {
+ axis = input_dims - 1;
+ output_dims = src_dims;
+ }
memory::format layout_type;
// In MKL, data format passed to mkl softmax op depends on dimension of
// the input tensor. Here "x" data format in MKL is used for 1 dim tensor,
// "nc" for 2 dim tensor, "tnc" for 3 dim tensor, "nchw" for 4 dim tensor,
- // and "ncdhw" for 5 dim tensor. Each of the simbols has the following
+ // and "ncdhw" for 5 dim tensor. Each of the symbols has the following
// meaning: n = batch, c = channels, t = sequence length, h = height, w =
- // width, d = depth
+ // width, d = depth. When src tensor is MKL, layout_type here is only used
+ // for setting TF layout type of output tensor. When input is TF Tensor,
+ // layout here is no special sense. We use axis to define on which
+ // dimension to do softmax.
switch (input_dims) {
case 1:
layout_type = memory::format::x;
@@ -81,13 +91,22 @@
layout_type = memory::format::tnc;
break;
case 4:
- layout_type = memory::format::nchw;
+ if (src_mkl_shape.IsMklTensor()) {
+ layout_type = memory::format::nhwc;
+ } else {
+ layout_type = memory::format::nchw;
+ }
break;
case 5:
- layout_type = memory::format::ncdhw;
+ if (src_mkl_shape.IsMklTensor()) {
+ layout_type = memory::format::ndhwc;
+ } else {
+ layout_type = memory::format::ncdhw;
+ }
break;
default:
- OP_REQUIRES_OK(context, errors::Aborted("Input dims must be <= 5 and >=1"));
+ OP_REQUIRES_OK(context,
+ errors::Aborted("Input dims must be <= 5 and >=1"));
return;
}
// Create softmax memory for src, dst: both are defined in mkl_util.h,
@@ -99,25 +118,17 @@
// construct input Tf layout. For TF layout, although input shape
// (src_dims) required is in MKL-DNN order, the layout is Tensorflow's
// layout
- auto src_md =
- src_mkl_shape.IsMklTensor()
- ? src_mkl_shape.GetMklLayout()
- : memory::desc(src_dims, MklDnnType<T>(), layout_type);
+ auto src_md = src_mkl_shape.IsMklTensor()
+ ? src_mkl_shape.GetMklLayout()
+ : memory::desc(src_dims, MklDnnType<T>(), layout_type);
- // src: setting memory descriptor and op memory descriptor
- // Basically following two functions maps the TF "src_tensor" to mkl
- // tensor object "src"
+ // src: setting memory descriptor
// following functions are in mkl_util.h
- // data format is "nc" for src and dst; since the src and dst buffer is
- // always in 2D shape
src.SetUsrMem(src_md, &src_tensor);
- src.SetOpMemDesc(src_dims, layout_type);
// creating a memory descriptor
- // passing outermost dim as default axis, where the softmax is applied
- int axis = input_dims - 1;
auto softmax_fwd_desc = softmax_forward::desc(prop_kind::forward_scoring,
- src.GetOpMemDesc(), axis);
+ src.GetUsrMemDesc(), axis);
auto softmax_fwd_pd =
softmax_forward::primitive_desc(softmax_fwd_desc, cpu_engine);
diff --git a/tensorflow/core/kernels/partitioned_function_ops.cc b/tensorflow/core/kernels/partitioned_function_ops.cc
index 30a8be1..fbecd90 100644
--- a/tensorflow/core/kernels/partitioned_function_ops.cc
+++ b/tensorflow/core/kernels/partitioned_function_ops.cc
@@ -184,12 +184,6 @@
OptimizationPassRegistry::Global()->RunGrouping(
OptimizationPassRegistry::POST_PLACEMENT, optimization_options),
done);
- OP_REQUIRES_OK_ASYNC(
- ctx,
- OptimizationPassRegistry::Global()->RunGrouping(
- OptimizationPassRegistry::POST_REWRITE_FOR_EXEC,
- optimization_options),
- done);
Device* cpu_device;
OP_REQUIRES_OK_ASYNC(
@@ -197,10 +191,19 @@
// Run grappler passes on the graph. It is possible that these are
// optimized by the graph executor already.
- OP_REQUIRES_OK_ASYNC(ctx,
- OptimizeGraph(ctx, fbody->ret_nodes, overlay_lib,
- device_set, cpu_device, &graph),
- done);
+ Status optimized = OptimizeGraph(ctx, fbody->ret_nodes, overlay_lib,
+ device_set, cpu_device, &graph);
+ if (!optimized.ok()) {
+ LOG(WARNING) << "Grappler optimization failed. Error: "
+ << optimized.error_message();
+ }
+
+ OP_REQUIRES_OK_ASYNC(
+ ctx,
+ OptimizationPassRegistry::Global()->RunGrouping(
+ OptimizationPassRegistry::POST_REWRITE_FOR_EXEC,
+ optimization_options),
+ done);
std::unordered_map<string, std::unique_ptr<Graph>> subgraphs;
OP_REQUIRES_OK_ASYNC(
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index 170b08b..4167b60 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -55,6 +55,7 @@
#include <vector>
#include "absl/strings/str_join.h"
+#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/resource_mgr.h"
@@ -84,6 +85,47 @@
OP_REQUIRES_OK(c, c->GetAttr("dtype", &dtype_));
}
+namespace {
+Status CopyVariable(int output_idx, OpKernelContext* ctx, const Tensor* t) {
+ Tensor* output;
+ Notification n;
+ Status status;
+ AllocatorAttributes attr;
+ if (t->dtype() == DT_VARIANT) {
+ attr.set_on_host(true);
+ }
+ TF_RETURN_IF_ERROR(
+ ctx->allocate_output(output_idx, t->shape(), &output, attr));
+ if (t->dtype() == DT_VARIANT) {
+ output->flat<Variant>() = t->flat<Variant>();
+ } else if (ctx->op_device_context() != nullptr) {
+ // TODO(apassos): remove the down_cast by just returning Device* from
+ // OpKernelContext
+ Device* device = static_cast<Device*>(ctx->device());
+ ctx->op_device_context()->CopyTensorInSameDevice(
+ t, device, output, [&n, &status](const Status& s) {
+ status = s;
+ n.Notify();
+ });
+ n.WaitForNotification();
+ return status;
+ } else {
+ switch (t->dtype()) {
+#define HANDLER(type) \
+ case DataTypeToEnum<type>::value: \
+ output->flat<type>() = t->flat<type>(); \
+ break;
+ TF_CALL_ALL_TYPES(HANDLER);
+#undef HANDLER
+ default:
+ return errors::Internal("Unsupported dtype", t->dtype());
+ }
+ }
+ return Status::OK();
+}
+
+} // namespace
+
void ReadVariableOp::Compute(OpKernelContext* ctx) {
Var* variable = nullptr;
const ResourceHandle& handle = HandleFromInput(ctx, 0);
@@ -100,12 +142,16 @@
// holding a shared lock to guarantee ordering of reads and
// writes.
tf_shared_lock ml(*variable->mu());
- const Tensor& t = *variable->tensor();
- OP_REQUIRES(ctx, dtype_ == t.dtype(),
+ const Tensor* t = variable->tensor();
+ OP_REQUIRES(ctx, dtype_ == t->dtype(),
errors::InvalidArgument(
"Trying to read variable with wrong dtype. Expected ",
- DataTypeString(dtype_), " got ", DataTypeString(t.dtype())));
- ctx->set_output(0, t);
+ DataTypeString(dtype_), " got ", DataTypeString(t->dtype())));
+ if (variable->copy_on_read_mode.load()) {
+ OP_REQUIRES_OK(ctx, CopyVariable(0, ctx, t));
+ } else {
+ ctx->set_output(0, *t);
+ }
}
ReadVariablesOp::ReadVariablesOp(OpKernelConstruction* c) : OpKernel(c) {
@@ -146,14 +192,18 @@
// holding a shared lock to guarantee ordering of reads and
// writes.
tf_shared_lock ml(*variables[i]->mu());
- const Tensor& t = *variables[i]->tensor();
- OP_REQUIRES(ctx, dtypes_[i] == t.dtype(),
+ OP_REQUIRES(ctx, dtypes_[i] == variables[i]->tensor()->dtype(),
errors::InvalidArgument(
"Trying to read variable ", handles[i]->name(),
" from Container: ", handles[i]->container(),
" with wrong dtype. Expected ", DataTypeString(dtypes_[i]),
- " got ", DataTypeString(t.dtype())));
- ctx->set_output(i, t);
+ " got ", DataTypeString(variables[i]->tensor()->dtype())));
+ if (variables[i]->copy_on_read_mode.load()) {
+ OP_REQUIRES_OK(ctx, CopyVariable(i, ctx, variables[i]->tensor()));
+ } else {
+ const Tensor& t = *variables[i]->tensor();
+ ctx->set_output(i, t);
+ }
}
}
@@ -308,8 +358,23 @@
"Trying to assign variable with wrong dtype. Expected ",
DataTypeString(variable->tensor()->dtype()), " got ",
DataTypeString(dtype_)));
+ if (variable->copy_on_read_mode.load()) {
+ PersistentTensor unused;
+ Tensor* tmp;
+ AllocatorAttributes attr;
+ attr.set_gpu_compatible(true);
+ attr.set_nic_compatible(true);
+ OP_REQUIRES_OK(context,
+ context->allocate_persistent(value.dtype(), value.shape(),
+ &unused, &tmp, attr));
+ functor::DenseUpdate<Device, T, ASSIGN> copy_functor;
+ copy_functor(context->eigen_device<Device>(), tmp->flat<T>(),
+ value.flat<T>());
+ *variable->tensor() = *tmp;
+ } else {
+ *variable->tensor() = value;
+ }
variable->is_initialized = true;
- *variable->tensor() = value;
}
private:
@@ -442,8 +507,9 @@
" using a Tensor with shape ",
value.shape().DebugString(),
", shapes must be equal."));
- OP_REQUIRES_OK(context,
- PrepareToUpdateVariable<Device, T>(context, var_tensor));
+ OP_REQUIRES_OK(
+ context, PrepareToUpdateVariable<Device, T>(
+ context, var_tensor, variable->copy_on_read_mode.load()));
functor::DenseUpdate<Device, T, Op> update_functor;
update_functor(context->eigen_device<Device>(), var_tensor->flat<T>(),
value.flat<T>());
@@ -524,6 +590,7 @@
Var* v = nullptr;
OP_REQUIRES_OK(c, LookupResource(c, HandleFromInput(c, 0), &v));
core::ScopedUnref su(v);
+ OP_REQUIRES_OK(c, EnsureSparseVariableAccess<Device, T>(c, v));
// NOTE: We hold the lock for the whole gather operation instead
// of increasing the reference count of v->tensor() to avoid a
// situation where a write to the same variable will see a
@@ -639,9 +706,9 @@
Var* v = nullptr;
OP_REQUIRES_OK(c, LookupResource(c, HandleFromInput(c, 0), &v));
core::ScopedUnref unref_v(v);
- mutex_lock ml(*v->mu());
+ OP_REQUIRES_OK(c, EnsureSparseVariableAccess<Device, T>(c, v));
+ tf_shared_lock ml(*v->mu());
Tensor* params = v->tensor();
- OP_REQUIRES_OK(c, PrepareToUpdateVariable<Device, T>(c, params));
const Tensor& indices = c->input(1);
const Tensor& updates = c->input(2);
diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc
index 63bb793..b466e57 100644
--- a/tensorflow/core/kernels/scatter_nd_op.cc
+++ b/tensorflow/core/kernels/scatter_nd_op.cc
@@ -231,6 +231,7 @@
Var* v;
OP_REQUIRES_OK(c, LookupResource(c, HandleFromInput(c, 0), &v));
core::ScopedUnref scoped_unref(v);
+ OP_REQUIRES_OK(c, EnsureSparseVariableAccess<Device, T>(c, v));
mutex_lock m(*v->mu());
DoCompute(c);
} else if (use_exclusive_lock_) {
@@ -258,7 +259,6 @@
Var* v;
OP_REQUIRES_OK(c, LookupResource(c, HandleFromInput(c, 0), &v));
Tensor* t = v->tensor();
- OP_REQUIRES_OK(c, PrepareToUpdateVariable<Device, T>(c, t));
params = *t;
params_shape = params.shape();
} else if (IsRefType(c->input_dtype(0))) {
diff --git a/tensorflow/core/kernels/scatter_op.cc b/tensorflow/core/kernels/scatter_op.cc
index 0fbde76..ee3c583 100644
--- a/tensorflow/core/kernels/scatter_op.cc
+++ b/tensorflow/core/kernels/scatter_op.cc
@@ -288,7 +288,7 @@
#define REGISTER_SCATTER_UPDATE_GPU(type) REGISTER_SCATTER_UPDATE(type, GPU);
TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SCATTER_ARITHMETIC_GPU);
-TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SCATTER_MINMAX_GPU);
+TF_CALL_GPU_NUMBER_TYPES(REGISTER_SCATTER_MINMAX_GPU);
TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SCATTER_UPDATE_GPU);
#endif // GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/scatter_op_gpu.cu.cc b/tensorflow/core/kernels/scatter_op_gpu.cu.cc
index 0df3293..d4defb8 100644
--- a/tensorflow/core/kernels/scatter_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/scatter_op_gpu.cu.cc
@@ -41,6 +41,7 @@
DEFINE_GPU_SPECS_INDEX(T, int32); \
DEFINE_GPU_SPECS_INDEX(T, int64);
+DEFINE_GPU_SPECS(Eigen::half);
DEFINE_GPU_SPECS(float);
DEFINE_GPU_SPECS(double);
// TODO: The following fails to compile.
diff --git a/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc b/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
index ac48202..a4e89f4 100644
--- a/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
+++ b/tensorflow/core/kernels/sparse_dense_binary_op_shared.cc
@@ -88,12 +88,12 @@
const auto rhs_dims = BCast::FromShape(dense_t->shape());
BCast b(lhs_dims, rhs_dims, false); // false for keeping the same num dims.
- // True iff (size(lhs) > size(rhs)), or (sizes equal, lhs cwise rhs).
+ // True iff (size(lhs) >= size(rhs)) and all dims in lhs is greater or equal
+ // to dims in rhs (from right to left).
auto VecGreaterEq = [](ArraySlice<int64> lhs, ArraySlice<int64> rhs) {
- if (lhs.size() > rhs.size()) return true;
if (lhs.size() < rhs.size()) return false;
- for (size_t i = 0; i < lhs.size(); ++i) {
- if (lhs[i] < rhs[i]) return false;
+ for (size_t i = 0; i < rhs.size(); ++i) {
+ if (lhs[lhs.size() - 1 - i] < rhs[rhs.size() - 1 - i]) return false;
}
return true;
};
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc
index 70a7ddb..6db68f9 100644
--- a/tensorflow/core/kernels/strided_slice_op.cc
+++ b/tensorflow/core/kernels/strided_slice_op.cc
@@ -307,9 +307,9 @@
OP_REQUIRES_OK(context,
LookupResource(context, HandleFromInput(context, 0), &v));
core::ScopedUnref scoped_unref(v);
- mutex_lock ml(*v->mu());
OP_REQUIRES_OK(context,
- PrepareToUpdateVariable<Device, T>(context, v->tensor()));
+ EnsureSparseVariableAccess<Device, T>(context, v));
+ mutex_lock ml(*v->mu());
old_lhs = v->tensor();
OP_REQUIRES(context, old_lhs->dtype() == DataTypeToEnum<T>::value,
errors::InvalidArgument(
diff --git a/tensorflow/core/kernels/summary_image_op.cc b/tensorflow/core/kernels/summary_image_op.cc
index 29b21ee..68f17c2 100644
--- a/tensorflow/core/kernels/summary_image_op.cc
+++ b/tensorflow/core/kernels/summary_image_op.cc
@@ -78,6 +78,11 @@
const int hw = h * w; // Compact these two dims for simplicity
const int depth = static_cast<int>(tensor.dim_size(3));
+ OP_REQUIRES(c, hw > 0 && depth > 0,
+ errors::InvalidArgument(
+ "input tensor must have non-zero dims. Found: [",
+ batch_size, ", ", h, ", ", w, ", ", depth, "]."));
+
Summary s;
if (tensor.dtype() == DT_UINT8) {
// For uint8 input, no normalization is necessary
diff --git a/tensorflow/core/kernels/tile_ops.cc b/tensorflow/core/kernels/tile_ops.cc
index d714876..b9b3761 100644
--- a/tensorflow/core/kernels/tile_ops.cc
+++ b/tensorflow/core/kernels/tile_ops.cc
@@ -325,6 +325,7 @@
TF_CALL_int16(HANDLE_TYPE_NAME);
TF_CALL_int64(HANDLE_TYPE_NAME);
TF_CALL_half(HANDLE_TYPE_NAME);
+ TF_CALL_bfloat16(HANDLE_TYPE_NAME);
TF_CALL_complex64(HANDLE_TYPE_NAME);
TF_CALL_complex128(HANDLE_TYPE_NAME);
diff --git a/tensorflow/core/kernels/training_op_helpers.cc b/tensorflow/core/kernels/training_op_helpers.cc
index 4262a54..20c08cf 100644
--- a/tensorflow/core/kernels/training_op_helpers.cc
+++ b/tensorflow/core/kernels/training_op_helpers.cc
@@ -19,70 +19,6 @@
namespace tensorflow {
-mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input,
- Var** maybe_resource) {
- *maybe_resource = nullptr;
- if (ctx->input_dtype(input) == DT_RESOURCE) {
- if (LookupResource(ctx, HandleFromInput(ctx, input), maybe_resource).ok()) {
- return (*maybe_resource)->mu();
- } else {
- ctx->CtxFailureWithWarning(
- errors::Internal("Invalid variable reference."));
- return nullptr;
- }
- }
- return ctx->input_ref_mutex(input);
-}
-
-// MaybeLockVariableInputMutexesInOrder is a helper function to acquire mutexes
-// in address order to mitigate deadlock. Returns a structure that, when
-// deleted, will release the acquired mutexes. Safe to pass duplicates - will
-// only lock each distinct mutex once. If do_lock is false, returns
-// immediately. Note that this silently doesn't lock mutexes for invalid
-// variable references; in all usages this is followed by GetInputTensor which
-// will signal a failure.
-VariableInputLockHolder MaybeLockVariableInputMutexesInOrder(
- OpKernelContext* ctx, bool do_lock, const std::vector<int>& input_ids) {
- bool any_resource = false;
- for (auto i : input_ids) {
- if (ctx->input_dtype(i) == DT_RESOURCE) {
- any_resource = true;
- break;
- }
- }
- if (!do_lock && !any_resource) {
- return VariableInputLockHolder({}, {});
- }
- std::vector<Var*> vars;
- std::vector<mutex*> mutexes;
- std::vector<int> acquire_order;
- for (auto input : input_ids) {
- Var* var;
- mutex* mutex = GetTrainingVariableMutex(ctx, input, &var);
- if (var) vars.push_back(var);
- // Only lock each mutex once if duplicates exist (n^2 but n is 2 or 3).
- if (std::find(mutexes.begin(), mutexes.end(), mutex) == mutexes.end()) {
- acquire_order.push_back(mutexes.size());
- mutexes.push_back(mutex);
- }
- }
- std::sort(acquire_order.begin(), acquire_order.end(),
- [&mutexes](int a, int b) { return mutexes[a] < mutexes[b]; });
-
- std::unique_ptr<std::vector<mutex_lock>> locks =
- MakeUnique<std::vector<mutex_lock>>();
- locks->reserve(acquire_order.size());
-
- for (auto input : acquire_order) {
- Var* var;
- mutex* mu = GetTrainingVariableMutex(ctx, input, &var);
- core::ScopedUnref scoped_unref(var);
- if (mu != nullptr) {
- locks->emplace_back(*mu);
- }
- }
- return VariableInputLockHolder(std::move(vars), std::move(locks));
-}
void MaybeForwardRefInputToRefOutput(OpKernelContext* ctx, int input,
int output) {
diff --git a/tensorflow/core/kernels/training_op_helpers.h b/tensorflow/core/kernels/training_op_helpers.h
index 9f173a8..715dd8a 100644
--- a/tensorflow/core/kernels/training_op_helpers.h
+++ b/tensorflow/core/kernels/training_op_helpers.h
@@ -17,30 +17,72 @@
#define TENSORFLOW_CORE_KERNELS_TRAINING_OP_HELPERS_H_
#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/variant_op_registry.h"
#include "tensorflow/core/kernels/dense_update_functor.h"
#include "tensorflow/core/kernels/variable_ops.h"
namespace tensorflow {
-// Returns a borrowed pointer to the mutex for the variable `input` in `ctx`.
-//
-// If `input` corresponds to a `DT_RESOURCE`-type variable input,
-// `*maybe_resource` will be updated to contain the underlying resource, and the
-// caller will be responsible for calling `Unref()` on that resource.
-mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input,
- Var** maybe_resource);
+// Must be called before performing a sparse operation on a variable. Ensures
+// that no concurrent dense operations can happen while holding the variable's
+// lock.
+template <typename Device, typename T>
+Status EnsureSparseVariableAccess(OpKernelContext* ctx, Var* var) {
+ if (var->copy_on_read_mode.load()) {
+ return Status::OK();
+ }
+ mutex_lock ml(*var->mu());
+ // Once copy-on-read mode is True the refcount is guaranteed to be 1. This can
+ // also happen if there are no concurrent reads of the variable and
+ // copy-on-read mode is false.
+ if (var->tensor()->RefCountIsOne()) {
+ var->copy_on_read_mode.store(true);
+ return Status::OK();
+ }
+ PersistentTensor unused;
+ Tensor* tmp;
+ if (std::is_same<T, Variant>::value) {
+ AllocatorAttributes attr;
+ attr.set_on_host(true);
+ TF_RETURN_IF_ERROR(ctx->allocate_persistent(
+ var->tensor()->dtype(), var->tensor()->shape(), &unused, &tmp, attr));
+
+ const auto elements_in = var->tensor()->flat<Variant>();
+ auto elements_out = tmp->flat<Variant>();
+ for (int64 i = 0; i < elements_in.size(); ++i) {
+ elements_out(i) = elements_in(i);
+ }
+ } else {
+ AllocatorAttributes attr;
+ attr.set_gpu_compatible(true);
+ attr.set_nic_compatible(true);
+ TF_RETURN_IF_ERROR(ctx->allocate_persistent(
+ var->tensor()->dtype(), var->tensor()->shape(), &unused, &tmp, attr));
+ functor::DenseUpdate<Device, T, ASSIGN> copy_functor;
+ copy_functor(ctx->eigen_device<Device>(), tmp->flat<T>(),
+ const_cast<const Tensor*>(var->tensor())->flat<T>());
+ }
+ *var->tensor() = *tmp;
+ var->copy_on_read_mode.store(true);
+ return Status::OK();
+}
// Utility structure that releases a sequence of borrowed mutexes when it is
// deleted.
struct VariableInputLockHolder {
public:
- VariableInputLockHolder(std::vector<Var*> vars,
- std::unique_ptr<std::vector<mutex_lock>> locks)
- : vars_(std::move(vars)), locks_(std::move(locks)) {}
+ VariableInputLockHolder(
+ std::vector<Var*> vars, std::unique_ptr<std::vector<mutex_lock>> locks,
+ std::unique_ptr<std::vector<tf_shared_lock>> shared_locks)
+ : vars_(std::move(vars)),
+ locks_(std::move(locks)),
+ shared_locks_(std::move(shared_locks)) {}
VariableInputLockHolder(VariableInputLockHolder&& other)
- : vars_(std::move(other.vars_)), locks_(std::move(other.locks_)) {}
+ : vars_(std::move(other.vars_)),
+ locks_(std::move(other.locks_)),
+ shared_locks_(std::move(other.shared_locks_)) {}
~VariableInputLockHolder() {
// Release the locks before unreffing the Vars, because each lock
@@ -56,10 +98,96 @@
// NOTE: Use a `std::unique_ptr` instead of moving in a vector directly,
// because a `std::vector<mutex_lock>` is not movable on all platforms.
std::unique_ptr<std::vector<mutex_lock>> locks_;
+ std::unique_ptr<std::vector<tf_shared_lock>> shared_locks_;
};
+// Returns a borrowed pointer to the mutex for the variable `input` in `ctx`.
+//
+// If `input` corresponds to a `DT_RESOURCE`-type variable input,
+// `*maybe_resource` will be updated to contain the underlying resource, and the
+// caller will be responsible for calling `Unref()` on that resource.
+template <typename Device, typename T>
+mutex* GetTrainingVariableMutex(OpKernelContext* ctx, int input, bool sparse,
+ Var** maybe_resource) {
+ *maybe_resource = nullptr;
+ if (ctx->input_dtype(input) == DT_RESOURCE) {
+ if (LookupResource(ctx, HandleFromInput(ctx, input), maybe_resource).ok()) {
+ if (sparse) {
+ EnsureSparseVariableAccess<Device, T>(ctx, *maybe_resource)
+ .IgnoreError();
+ }
+ return (*maybe_resource)->mu();
+ } else {
+ ctx->CtxFailureWithWarning(
+ errors::Internal("Invalid variable reference."));
+ return nullptr;
+ }
+ }
+ return ctx->input_ref_mutex(input);
+}
+
+// MaybeLockVariableInputMutexesInOrder is a helper function to acquire mutexes
+// in address order to mitigate deadlock. Returns a structure that, when
+// deleted, will release the acquired mutexes. Safe to pass duplicates - will
+// only lock each distinct mutex once. If sparse is true will ensure the
+// variable gets switched to copy-on-read mode before trying to acquire the
+// locks. If do_lock is false, returns immediately for reference variables. For
+// resource variables in copy-on-read-mode it will grab a shared lock if do_lock
+// is false, exclusive lock otherwise. Note that this silently doesn't lock
+// mutexes for invalid variable references; in all usages this is followed by
+// GetInputTensor which will signal a failure.
+template <typename Device, typename T>
VariableInputLockHolder MaybeLockVariableInputMutexesInOrder(
- OpKernelContext* ctx, bool do_lock, const std::vector<int>& input_ids);
+ OpKernelContext* ctx, bool do_lock, bool sparse,
+ const std::vector<int>& input_ids) {
+ bool any_resource = false;
+ for (auto i : input_ids) {
+ if (ctx->input_dtype(i) == DT_RESOURCE) {
+ any_resource = true;
+ break;
+ }
+ }
+ if (!do_lock && !any_resource) {
+ return VariableInputLockHolder({}, {}, {});
+ }
+ std::vector<Var*> vars;
+ std::vector<mutex*> mutexes;
+ std::vector<int> acquire_order;
+ for (auto input : input_ids) {
+ Var* var;
+ mutex* mutex =
+ GetTrainingVariableMutex<Device, T>(ctx, input, sparse, &var);
+ if (var) vars.push_back(var);
+ // Only lock each mutex once if duplicates exist (n^2 but n is 2 or 3).
+ if (std::find(mutexes.begin(), mutexes.end(), mutex) == mutexes.end()) {
+ acquire_order.push_back(mutexes.size());
+ mutexes.push_back(mutex);
+ }
+ }
+ std::sort(acquire_order.begin(), acquire_order.end(),
+ [&mutexes](int a, int b) { return mutexes[a] < mutexes[b]; });
+
+ std::unique_ptr<std::vector<mutex_lock>> locks =
+ absl::make_unique<std::vector<mutex_lock>>();
+ std::unique_ptr<std::vector<tf_shared_lock>> shared_locks =
+ absl::make_unique<std::vector<tf_shared_lock>>();
+ locks->reserve(acquire_order.size());
+
+ for (auto input : acquire_order) {
+ Var* var;
+ mutex* mu = GetTrainingVariableMutex<Device, T>(ctx, input, sparse, &var);
+ core::ScopedUnref scoped_unref(var);
+ if (mu != nullptr) {
+ if (!sparse || do_lock) {
+ locks->emplace_back(*mu);
+ } else {
+ shared_locks->emplace_back(*mu);
+ }
+ }
+ }
+ return VariableInputLockHolder(std::move(vars), std::move(locks),
+ std::move(shared_locks));
+}
void MaybeForwardRefInputToRefOutput(OpKernelContext* ctx, int input,
int output);
@@ -68,8 +196,9 @@
// reference count of 1 before you update it.
// REQUIRES: If you pass in variable->tensor(), *variable->mu() must be held.
template <typename Device, typename T>
-Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor) {
- if (!tensor->RefCountIsOne()) {
+Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor,
+ bool copy_on_read_mode) {
+ if (copy_on_read_mode || !tensor->RefCountIsOne()) {
// Tensor's buffer is in use by some read, so we need to copy before
// updating.
PersistentTensor unused;
@@ -100,12 +229,14 @@
return Status::OK();
}
-// This gives you `*out`, a tensor you can update, corresponding to a
-// variable passed as input index `input`. This handles the
-// differences between reference and resource variables. For resource
-// variables, we ensure `*out` has a reference count of 1 (using
-// PrepareToUpdateVariable() to copy if necessary) unless
-// sparse && !lock_held, in which case it never copies.
+// This gives you `*out`, a tensor you can update, corresponding to a variable
+// passed as input index `input`. This handles the differences between
+// reference and resource variables. For reference variables we can just grab
+// the tensor, grabbing the lock if lock_held is False.
+//
+// For resource variables we, if sparse is true, ensure it's in copy-on-read
+// mode, and then, regardless of the value of sparse, ensure its refcount is 1
+// (by potentially copying its contents). In this case lock_held is ignored.
template <typename Device, typename T>
Status GetInputTensorFromVariable(OpKernelContext* ctx, int input,
bool lock_held, bool sparse, Tensor* out) {
@@ -113,7 +244,13 @@
Var* var;
TF_RETURN_IF_ERROR(LookupResource(ctx, HandleFromInput(ctx, input), &var));
core::ScopedUnref unref_var(var);
- TF_RETURN_IF_ERROR(PrepareToUpdateVariable<Device, T>(ctx, var->tensor()));
+ if (sparse) {
+ TF_RETURN_IF_ERROR(EnsureSparseVariableAccess<Device, T>(ctx, var));
+ *out = *var->tensor();
+ return Status::OK();
+ }
+ TF_RETURN_IF_ERROR(PrepareToUpdateVariable<Device, T>(
+ ctx, var->tensor(), var->copy_on_read_mode.load()));
*out = *var->tensor();
return Status::OK();
}
diff --git a/tensorflow/core/kernels/training_ops.cc b/tensorflow/core/kernels/training_ops.cc
index 6504ad1..b2239ab 100644
--- a/tensorflow/core/kernels/training_ops.cc
+++ b/tensorflow/core/kernels/training_ops.cc
@@ -465,11 +465,12 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -506,11 +507,12 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<SYCLDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<SYCLDevice, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -600,7 +602,8 @@
void Compute(OpKernelContext* ctx) override {
Var* resource;
- mutex* mu = GetTrainingVariableMutex(ctx, 0, &resource);
+ const bool sparse = false;
+ mutex* mu = GetTrainingVariableMutex<Device, T>(ctx, 0, sparse, &resource);
core::ScopedUnref scoped_unref(resource);
if (use_exclusive_lock_ && mu != nullptr) {
mutex_lock l1(*mu);
@@ -624,14 +627,16 @@
void DoValidate(OpKernelContext* ctx) {
Tensor var;
+ const bool sparse = false;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
Tensor accum_update;
- OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &accum_update));
+ OP_REQUIRES_OK(
+ ctx, GetInputTensorFromVariable<Device, T>(ctx, 2, use_exclusive_lock_,
+ sparse, &accum_update));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -678,14 +683,16 @@
void DoCompute(OpKernelContext* ctx) {
const Device& device = ctx->template eigen_device<Device>();
Tensor var;
+ const bool sparse = false;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
Tensor accum_update;
- OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &accum_update));
+ OP_REQUIRES_OK(
+ ctx, GetInputTensorFromVariable<Device, T>(ctx, 2, use_exclusive_lock_,
+ sparse, &accum_update));
const Tensor& lr = ctx->input(3);
const Tensor& rho = ctx->input(4);
@@ -751,7 +758,8 @@
void Compute(OpKernelContext* ctx) override {
Var* var;
- mutex* mu = GetTrainingVariableMutex(ctx, 0, &var);
+ const bool sparse = true;
+ mutex* mu = GetTrainingVariableMutex<CPUDevice, T>(ctx, 0, sparse, &var);
core::ScopedUnref scoped_unref(var);
// mu_accum is actually the same mutex as mu_var since currently we use a
// global mutex.
@@ -767,14 +775,16 @@
void DoCompute(OpKernelContext* ctx) {
Tensor var;
+ const bool sparse = true;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum_grad;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &accum_grad));
+ ctx, 1, use_exclusive_lock_, sparse, &accum_grad));
Tensor accum_update;
- OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 2, use_exclusive_lock_, true, &accum_update));
+ OP_REQUIRES_OK(ctx,
+ GetInputTensorFromVariable<CPUDevice, T>(
+ ctx, 2, use_exclusive_lock_, sparse, &accum_update));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -907,11 +917,12 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -976,11 +987,12 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
OP_REQUIRES(ctx, TensorShapeUtils::IsVectorOrHigher(var.shape()),
errors::InvalidArgument("var must be at least 1 dimensional"));
@@ -1121,14 +1133,15 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -1214,14 +1227,15 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -1316,14 +1330,15 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -1456,14 +1471,15 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -1628,19 +1644,20 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor gradient_accum;
OP_REQUIRES_OK(
ctx, GetInputTensorFromVariable<Device, T>(ctx, 1, use_exclusive_lock_,
- false, &gradient_accum));
+ sparse, &gradient_accum));
Tensor gradient_squared_accum;
OP_REQUIRES_OK(
ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &gradient_squared_accum));
+ ctx, 2, use_exclusive_lock_, sparse, &gradient_squared_accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -1729,19 +1746,20 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor gradient_accum;
OP_REQUIRES_OK(ctx,
GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &gradient_accum));
+ ctx, 1, use_exclusive_lock_, sparse, &gradient_accum));
Tensor gradient_squared_accum;
OP_REQUIRES_OK(
ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 2, use_exclusive_lock_, true, &gradient_squared_accum));
+ ctx, 2, use_exclusive_lock_, sparse, &gradient_squared_accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -1927,18 +1945,19 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
Tensor linear;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &linear));
+ ctx, 2, use_exclusive_lock_, sparse, &linear));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2079,17 +2098,18 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, true, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
Tensor linear;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, true, &linear));
+ ctx, 2, use_exclusive_lock_, sparse, &linear));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2353,15 +2373,16 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2454,15 +2475,16 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2572,15 +2594,16 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2671,15 +2694,16 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor accum;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &accum));
+ ctx, 1, use_exclusive_lock_, sparse, &accum));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2783,18 +2807,19 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor m;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &m));
+ ctx, 1, use_exclusive_lock_, sparse, &m));
Tensor v;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &v));
+ ctx, 2, use_exclusive_lock_, sparse, &v));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -2873,18 +2898,19 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<SYCLDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<SYCLDevice, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor m;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<SYCLDevice, T>(
- ctx, 1, use_exclusive_lock_, false, &m));
+ ctx, 1, use_exclusive_lock_, sparse, &m));
Tensor v;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<SYCLDevice, T>(
- ctx, 2, use_exclusive_lock_, false, &v));
+ ctx, 2, use_exclusive_lock_, sparse, &v));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -3043,21 +3069,22 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor m;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &m));
+ ctx, 1, use_exclusive_lock_, sparse, &m));
Tensor v;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &v));
+ ctx, 2, use_exclusive_lock_, sparse, &v));
Tensor vhat;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 3, use_exclusive_lock_, false, &vhat));
+ ctx, 3, use_exclusive_lock_, sparse, &vhat));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -3184,18 +3211,19 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor m;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &m));
+ ctx, 1, use_exclusive_lock_, sparse, &m));
Tensor v;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &v));
+ ctx, 2, use_exclusive_lock_, sparse, &v));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -3312,18 +3340,19 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor ms;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &ms));
+ ctx, 1, use_exclusive_lock_, sparse, &ms));
Tensor mom;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &mom));
+ ctx, 2, use_exclusive_lock_, sparse, &mom));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -3394,21 +3423,22 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2, 3});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2, 3});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor mg;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &mg));
+ ctx, 1, use_exclusive_lock_, sparse, &mg));
Tensor ms;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 2, use_exclusive_lock_, false, &ms));
+ ctx, 2, use_exclusive_lock_, sparse, &ms));
Tensor mom;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 3, use_exclusive_lock_, false, &mom));
+ ctx, 3, use_exclusive_lock_, sparse, &mom));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -3553,18 +3583,19 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor ms;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &ms));
+ ctx, 1, use_exclusive_lock_, sparse, &ms));
Tensor mom;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 2, use_exclusive_lock_, true, &mom));
+ ctx, 2, use_exclusive_lock_, sparse, &mom));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -3682,21 +3713,22 @@
}
void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
- auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
- {0, 1, 2, 3});
+ const bool sparse = true;
+ auto locks = MaybeLockVariableInputMutexesInOrder<CPUDevice, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1, 2, 3});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 0, use_exclusive_lock_, true, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor mg;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 1, use_exclusive_lock_, true, &mg));
+ ctx, 1, use_exclusive_lock_, sparse, &mg));
Tensor ms;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 2, use_exclusive_lock_, true, &ms));
+ ctx, 2, use_exclusive_lock_, sparse, &ms));
Tensor mom;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<CPUDevice, T>(
- ctx, 3, use_exclusive_lock_, true, &mom));
+ ctx, 3, use_exclusive_lock_, sparse, &mom));
OP_REQUIRES(
ctx, var.IsInitialized(),
@@ -3852,15 +3884,16 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor m;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &m));
+ ctx, 1, use_exclusive_lock_, sparse, &m));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
@@ -3958,15 +3991,16 @@
}
void Compute(OpKernelContext* ctx) override {
- auto locks =
- MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+ const bool sparse = false;
+ auto locks = MaybeLockVariableInputMutexesInOrder<Device, T>(
+ ctx, use_exclusive_lock_, sparse, {0, 1});
Tensor var;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 0, use_exclusive_lock_, false, &var));
+ ctx, 0, use_exclusive_lock_, sparse, &var));
Tensor m;
OP_REQUIRES_OK(ctx, GetInputTensorFromVariable<Device, T>(
- ctx, 1, use_exclusive_lock_, false, &m));
+ ctx, 1, use_exclusive_lock_, sparse, &m));
OP_REQUIRES(
ctx, var.IsInitialized(),
errors::FailedPrecondition(
diff --git a/tensorflow/core/kernels/training_ops_test.cc b/tensorflow/core/kernels/training_ops_test.cc
index 1ec57b4..2dcc4a5 100644
--- a/tensorflow/core/kernels/training_ops_test.cc
+++ b/tensorflow/core/kernels/training_ops_test.cc
@@ -151,40 +151,6 @@
}
BENCHMARK(BM_Momentum)->Arg(128 << 10)->Arg(256 << 10);
-static void KerasMomentum(int32 n, Graph** init_g, Graph** train_g) {
- TensorShape shape({n});
- {
- Graph* g = new Graph(OpRegistry::Global());
- auto var = Var(g, n);
- auto accum = Var(g, n);
- auto zero = Zeros(g, n);
- test::graph::Assign(g, var, zero);
- test::graph::Assign(g, accum, zero);
- *init_g = g;
- }
- {
- Graph* g = new Graph(OpRegistry::Global());
- auto var = Var(g, n);
- auto accum = Var(g, n);
- auto lr = Scalar(g, 0.01);
- auto grad = Random(g, n);
- auto mom = Scalar(g, 0.01);
- test::graph::Multi(g, "ApplyKerasMomentum", {var, accum, lr, grad, mom});
- *train_g = g;
- }
-}
-
-static void BM_KerasMomentum(int iters, int params) {
- const int64 tot = static_cast<int64>(iters) * params;
- testing::ItemsProcessed(tot);
- testing::BytesProcessed(tot * sizeof(float));
- Graph* init;
- Graph* train;
- KerasMomentum(params, &init, &train);
- test::Benchmark("cpu", train, GetOptions(), init).Run(iters);
-}
-BENCHMARK(BM_KerasMomentum)->Arg(128 << 10)->Arg(256 << 10);
-
static void Adam(int32 n, Graph** init_g, Graph** train_g) {
TensorShape shape({n});
{
@@ -228,50 +194,6 @@
}
BENCHMARK(BM_Adam)->Arg(128 << 10)->Arg(256 << 10);
-static void AdamWithAmsgrad(int32 n, Graph** init_g, Graph** train_g) {
- TensorShape shape({n});
- {
- Graph* g = new Graph(OpRegistry::Global());
- auto var = Var(g, n);
- auto m = Var(g, n);
- auto v = Var(g, n);
- auto zero = Zeros(g, n);
- test::graph::Assign(g, var, zero);
- test::graph::Assign(g, m, zero);
- test::graph::Assign(g, v, zero);
- *init_g = g;
- }
- {
- Graph* g = new Graph(OpRegistry::Global());
- auto var = Var(g, n);
- auto m = Var(g, n);
- auto v = Var(g, n);
- auto vhat = Var(g, n);
- auto beta1_power = Scalar(g, 0.9);
- auto beta2_power = Scalar(g, 0.99);
- auto lr = Scalar(g, 0.01);
- auto beta1 = Scalar(g, 0.9);
- auto beta2 = Scalar(g, 0.99);
- auto epsilon = Scalar(g, 1e-8);
- auto grad = Random(g, n);
- test::graph::Multi(g, "ApplyAdamWithAmsgrad",
- {var, m, v, vhat, beta1_power, beta2_power, lr, beta1,
- beta2, epsilon, grad});
- *train_g = g;
- }
-}
-
-static void BM_AdamWithAmsgrad(int iters, int params) {
- const int64 tot = static_cast<int64>(iters) * params;
- testing::ItemsProcessed(tot);
- testing::BytesProcessed(tot * sizeof(float));
- Graph* init;
- Graph* train;
- AdamWithAmsgrad(params, &init, &train);
- test::Benchmark("cpu", train, GetOptions(), init).Run(iters);
-}
-BENCHMARK(BM_AdamWithAmsgrad)->Arg(128 << 10)->Arg(256 << 10);
-
static void RMSProp(int32 n, Graph** init_g, Graph** train_g) {
TensorShape shape({n});
{
diff --git a/tensorflow/core/kernels/unicode_ops.cc b/tensorflow/core/kernels/unicode_ops.cc
index 6c4ed1e..c9c2ac1 100644
--- a/tensorflow/core/kernels/unicode_ops.cc
+++ b/tensorflow/core/kernels/unicode_ops.cc
@@ -350,10 +350,10 @@
REGISTER_KERNEL_BUILDER(Name("UnicodeTranscode").Device(DEVICE_CPU),
UnicodeTranscodeOp);
-class UnicodeDecodeWithOffsetsOp : public OpKernel {
+class UnicodeDecodeBaseOp : public OpKernel {
public:
- explicit UnicodeDecodeWithOffsetsOp(OpKernelConstruction* ctx)
- : OpKernel(ctx) {
+ explicit UnicodeDecodeBaseOp(OpKernelConstruction* ctx, bool generate_offsets)
+ : OpKernel(ctx), generate_offsets_(generate_offsets) {
OP_REQUIRES_OK(ctx, GetErrorOptions(ctx, &error_options_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("input_encoding", &input_encoding_));
// Make a temporary UConverter to ensure it will create without error
@@ -369,7 +369,7 @@
}
void Decode(OpKernelContext* ctx, std::vector<UChar32>* char_values,
- std::vector<int64>* offset_values, int* string_length,
+ std::vector<int64>* offset_values, int* current_offset,
int64* next_row_split, UChar32 char_value, int char_length,
bool found_any_format_error) {
if (error_options_.error_on_malformatting && found_any_format_error) {
@@ -379,7 +379,8 @@
UChar32 decoded_value = char_value;
if (ShouldHandleFormatError(error_options_, char_value,
found_any_format_error)) {
- if (error_options_.elide_replacement) {
+ if (error_options_.elide_replacement && (offset_values != nullptr)) {
+ *current_offset += char_length;
return;
} else {
decoded_value = error_options_.subst;
@@ -390,8 +391,10 @@
char_values->push_back(decoded_value);
// Emit the byte offset
- offset_values->push_back(*string_length);
- *string_length += char_length;
+ if (offset_values != nullptr) {
+ offset_values->push_back(*current_offset);
+ *current_offset += char_length;
+ }
*next_row_split += 1;
}
@@ -428,42 +431,63 @@
// the fields needed to construct a RaggedTensor.
out_row_splits(row_split_index) = next_row_split;
row_split_index++;
- int string_length = 0;
+ int current_offset = 0;
IterateUnicodeString(
input, input_encoder->converter_,
- std::bind(&UnicodeDecodeWithOffsetsOp::Decode, this, ctx,
- &char_values, &offset_values, &string_length,
- &next_row_split, std::placeholders::_1,
- std::placeholders::_2, std::placeholders::_3));
+ std::bind(&UnicodeDecodeBaseOp::Decode, this, ctx, &char_values,
+ &offset_values, ¤t_offset, &next_row_split,
+ std::placeholders::_1, std::placeholders::_2,
+ std::placeholders::_3));
}
out_row_splits(row_split_index) = next_row_split;
- DCHECK(offset_values.size() == char_values.size());
Tensor* output_char_values;
OP_REQUIRES_OK(
ctx, ctx->allocate_output("char_values",
{static_cast<int64>(char_values.size())},
&output_char_values));
- Tensor* output_offset_values;
- OP_REQUIRES_OK(
- ctx, ctx->allocate_output("char_to_byte_starts",
- {static_cast<int64>(offset_values.size())},
- &output_offset_values));
auto out_char_values = output_char_values->vec<int32>();
- auto out_offset_values = output_offset_values->vec<int64>();
+ if (generate_offsets_) {
+ DCHECK(offset_values.size() == char_values.size());
+ Tensor* output_offset_values;
+ OP_REQUIRES_OK(
+ ctx, ctx->allocate_output("char_to_byte_starts",
+ {static_cast<int64>(offset_values.size())},
+ &output_offset_values));
+ auto out_offset_values = output_offset_values->vec<int64>();
- // Load output tensors from intermediate value arrays.
- for (int i = 0; i < char_values.size(); ++i) {
- out_char_values(i) = static_cast<int32>(char_values[i]);
- out_offset_values(i) = offset_values[i];
+ // Load output tensors from intermediate value arrays.
+ for (int i = 0; i < char_values.size(); ++i) {
+ out_char_values(i) = static_cast<int32>(char_values[i]);
+ out_offset_values(i) = offset_values[i];
+ }
+ } else {
+ for (int i = 0; i < char_values.size(); ++i) {
+ out_char_values(i) = static_cast<int32>(char_values[i]);
+ }
}
}
private:
string input_encoding_;
ErrorOptions error_options_;
+ bool generate_offsets_ = false;
};
+class UnicodeDecodeOp : public UnicodeDecodeBaseOp {
+ public:
+ explicit UnicodeDecodeOp(OpKernelConstruction* ctx)
+ : UnicodeDecodeBaseOp(ctx, false) {}
+};
+
+class UnicodeDecodeWithOffsetsOp : public UnicodeDecodeBaseOp {
+ public:
+ explicit UnicodeDecodeWithOffsetsOp(OpKernelConstruction* ctx)
+ : UnicodeDecodeBaseOp(ctx, true) {}
+};
+
+REGISTER_KERNEL_BUILDER(Name("UnicodeDecode").Device(DEVICE_CPU),
+ UnicodeDecodeOp);
REGISTER_KERNEL_BUILDER(Name("UnicodeDecodeWithOffsets").Device(DEVICE_CPU),
UnicodeDecodeWithOffsetsOp);
@@ -493,7 +517,7 @@
const Tensor& input_splits = context->input(1);
const auto input_splits_flat = input_splits.flat<int64>();
- // Since we limit to a 2-D input (inner_values of rank 1 and a single splits
+ // Since we limit to a 2-D input (flat_values of rank 1 and a single splits
// tensor), our output dimension will be 1 with it's size equal to the
// number of splits (outer dimension or ragged tensor).
TensorShape output_shape({input_splits.dim_size(0) - 1});
diff --git a/tensorflow/core/nccl/nccl_manager_test.cc b/tensorflow/core/nccl/nccl_manager_test.cc
index f43103e..f9ed4d0 100644
--- a/tensorflow/core/nccl/nccl_manager_test.cc
+++ b/tensorflow/core/nccl/nccl_manager_test.cc
@@ -65,6 +65,7 @@
static void SetUpTestCase() {
setenv("NCCL_DEBUG", "WARN", 1 /* replace */);
+ setenv("NCCL_LAUNCH_MODE", "PARALLEL", 1 /* replace */);
devices_ = new std::vector<std::unique_ptr<BaseGPUDevice>>(GetGPUDevices());
LOG(ERROR) << "Running test with " << devices_->size() << " gpus";
}
@@ -200,7 +201,7 @@
// Test basic sum reduction.
TYPED_TEST(NcclManagerTest, BasicSumReduction) {
- const int num_ranks = this->NumGPUs();
+ const int num_ranks = 4;
for (int op = 0; op < 4; ++op) {
ncclRedOp_t reduction_op = static_cast<ncclRedOp_t>(op);
@@ -208,6 +209,7 @@
this->MakeTestCase(num_ranks, reduction_op, TensorShape({2, 3}), 0.0f));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(rank);
+ VLOG(2) << "rank " << rank << " device " << device->name();
auto* event_mgr = device->tensorflow_gpu_device_info()->event_mgr;
auto* stream = device->tensorflow_gpu_device_info()->stream;
NcclManager::instance()->AddToAllReduce(
@@ -224,14 +226,12 @@
// Same as the Basic test, but with multiple threads launching parts of many
// reductions.
//
-// Testing the multi-rank execution is currently reduced as it can hang when run
-// with num_ranks > devices->size(), for some GPUs (e.g. K20m).
-// To test the higher settings, increase num_ranks,
-// num_collectives_per_iteration and time_limit_micros.
+// To run test longer, increase num_ranks, num_collectives_per_iteration and
+// time_limit_micros.
TYPED_TEST(NcclManagerTest, MultipleCallers) {
- const int num_ranks = this->NumGPUs();
+ const int num_ranks = 4;
const int num_collectives_per_iteration = 10; // 1000;
- const int num_threads = 3;
+ const int num_threads = num_ranks * 2;
const int time_limit_micros = 100; // 60 * 30 * 1000 * 1000;
int64 start = Env::Default()->NowMicros();
diff --git a/tensorflow/core/ops/compat/ops_history.v1.pbtxt b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
index 8022c39..602d4a0 100644
--- a/tensorflow/core/ops/compat/ops_history.v1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
@@ -21596,6 +21596,17 @@
is_stateful: true
}
op {
+ name: "ExperimentalDatasetCardinality"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "cardinality"
+ type: DT_INT64
+ }
+}
+op {
name: "ExperimentalDatasetToTFRecord"
input_arg {
name: "input_dataset"
@@ -22095,6 +22106,61 @@
}
}
op {
+ name: "ExperimentalMapAndBatchDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "other_arguments"
+ type_list_attr: "Targuments"
+ }
+ input_arg {
+ name: "batch_size"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "num_parallel_calls"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "drop_remainder"
+ type: DT_BOOL
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "Targuments"
+ type: "list(type)"
+ has_minimum: true
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "ExperimentalMapDataset"
input_arg {
name: "input_dataset"
@@ -22138,6 +22204,56 @@
}
}
op {
+ name: "ExperimentalMapDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "other_arguments"
+ type_list_attr: "Targuments"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "Targuments"
+ type: "list(type)"
+ has_minimum: true
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "use_inter_op_parallelism"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "ExperimentalMatchingFilesDataset"
input_arg {
name: "patterns"
@@ -22276,6 +22392,61 @@
}
}
op {
+ name: "ExperimentalNumaMapAndBatchDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "other_arguments"
+ type_list_attr: "Targuments"
+ }
+ input_arg {
+ name: "batch_size"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "num_parallel_calls"
+ type: DT_INT64
+ }
+ input_arg {
+ name: "drop_remainder"
+ type: DT_BOOL
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "Targuments"
+ type: "list(type)"
+ has_minimum: true
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "ExperimentalParallelInterleaveDataset"
input_arg {
name: "input_dataset"
@@ -22580,6 +22751,59 @@
}
}
op {
+ name: "ExperimentalScanDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "initial_state"
+ type_list_attr: "Tstate"
+ }
+ input_arg {
+ name: "other_arguments"
+ type_list_attr: "Targuments"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "Tstate"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "Targuments"
+ type: "list(type)"
+ has_minimum: true
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "ExperimentalSetStatsAggregatorDataset"
input_arg {
name: "input_dataset"
@@ -31138,6 +31362,46 @@
}
}
op {
+ name: "Lu"
+ input_arg {
+ name: "input"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "lu"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "p"
+ type_attr: "output_idx_type"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_DOUBLE
+ type: DT_FLOAT
+ type: DT_COMPLEX64
+ type: DT_COMPLEX128
+ }
+ }
+ }
+ attr {
+ name: "output_idx_type"
+ type: "type"
+ default_value {
+ type: DT_INT32
+ }
+ allowed_values {
+ list {
+ type: DT_INT32
+ type: DT_INT64
+ }
+ }
+ }
+}
+op {
name: "MakeIterator"
input_arg {
name: "dataset"
@@ -31304,6 +31568,56 @@
}
}
op {
+ name: "MapDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "other_arguments"
+ type_list_attr: "Targuments"
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "Targuments"
+ type: "list(type)"
+ has_minimum: true
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "use_inter_op_parallelism"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "MapDefun"
input_arg {
name: "arguments"
@@ -39703,6 +40017,67 @@
}
}
op {
+ name: "ParallelMapDataset"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ input_arg {
+ name: "other_arguments"
+ type_list_attr: "Targuments"
+ }
+ input_arg {
+ name: "num_parallel_calls"
+ type: DT_INT32
+ }
+ output_arg {
+ name: "handle"
+ type: DT_VARIANT
+ }
+ attr {
+ name: "f"
+ type: "func"
+ }
+ attr {
+ name: "Targuments"
+ type: "list(type)"
+ has_minimum: true
+ }
+ attr {
+ name: "output_types"
+ type: "list(type)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "output_shapes"
+ type: "list(shape)"
+ has_minimum: true
+ minimum: 1
+ }
+ attr {
+ name: "use_inter_op_parallelism"
+ type: "bool"
+ default_value {
+ b: true
+ }
+ }
+ attr {
+ name: "sloppy"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "ParameterizedTruncatedNormal"
input_arg {
name: "shape"
@@ -77433,6 +77808,53 @@
}
}
op {
+ name: "UnicodeDecode"
+ input_arg {
+ name: "input"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "row_splits"
+ type: DT_INT64
+ }
+ output_arg {
+ name: "char_values"
+ type: DT_INT32
+ }
+ attr {
+ name: "input_encoding"
+ type: "string"
+ }
+ attr {
+ name: "errors"
+ type: "string"
+ default_value {
+ s: "replace"
+ }
+ allowed_values {
+ list {
+ s: "strict"
+ s: "replace"
+ s: "ignore"
+ }
+ }
+ }
+ attr {
+ name: "replacement_char"
+ type: "int"
+ default_value {
+ i: 65533
+ }
+ }
+ attr {
+ name: "replace_control_characters"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "UnicodeDecodeWithOffsets"
input_arg {
name: "input"
@@ -78688,6 +79110,17 @@
is_stateful: true
}
op {
+ name: "UnwrapDatasetVariant"
+ input_arg {
+ name: "input_handle"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "output_handle"
+ type: DT_VARIANT
+ }
+}
+op {
name: "UpperBound"
input_arg {
name: "sorted_inputs"
@@ -79121,6 +79554,17 @@
}
}
op {
+ name: "WrapDatasetVariant"
+ input_arg {
+ name: "input_handle"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "output_handle"
+ type: DT_VARIANT
+ }
+}
+op {
name: "WriteAudioSummary"
input_arg {
name: "writer"
diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc
index b2ea637..1c11716 100644
--- a/tensorflow/core/ops/dataset_ops.cc
+++ b/tensorflow/core/ops/dataset_ops.cc
@@ -144,6 +144,7 @@
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
.Attr("use_inter_op_parallelism: bool = true")
+ .Attr("preserve_cardinality: bool = false")
.SetShapeFn(shape_inference::ScalarShape);
REGISTER_OP("ParallelMapDataset")
@@ -157,6 +158,7 @@
.Attr("output_shapes: list(shape) >= 1")
.Attr("use_inter_op_parallelism: bool = true")
.Attr("sloppy: bool = false")
+ .Attr("preserve_cardinality: bool = false")
.SetShapeFn(shape_inference::ScalarShape);
REGISTER_OP("PrefetchDataset")
@@ -685,6 +687,16 @@
return Status::OK();
});
+REGISTER_OP("WrapDatasetVariant")
+ .Input("input_handle: variant")
+ .Output("output_handle: variant")
+ .SetShapeFn(shape_inference::ScalarShape);
+
+REGISTER_OP("UnwrapDatasetVariant")
+ .Input("input_handle: variant")
+ .Output("output_handle: variant")
+ .SetShapeFn(shape_inference::ScalarShape);
+
REGISTER_OP("MultiDeviceIterator")
.Output("handle: resource")
.Attr("devices: list(string) >= 1")
diff --git a/tensorflow/core/ops/experimental_dataset_ops.cc b/tensorflow/core/ops/experimental_dataset_ops.cc
index 0711034..f904e25 100644
--- a/tensorflow/core/ops/experimental_dataset_ops.cc
+++ b/tensorflow/core/ops/experimental_dataset_ops.cc
@@ -174,6 +174,7 @@
.Attr("Targuments: list(type) >= 0")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
+ .Attr("preserve_cardinality: bool = false")
.SetShapeFn([](shape_inference::InferenceContext* c) {
// Use index from the end to retrieve the Input shapes,
// so that to avoid guessing the length of "other_arguments".
@@ -198,6 +199,7 @@
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
.Attr("use_inter_op_parallelism: bool = true")
+ .Attr("preserve_cardinality: bool = false")
.SetShapeFn(shape_inference::ScalarShape);
REGISTER_OP("ExperimentalMatchingFilesDataset")
@@ -277,6 +279,7 @@
.Attr("Targuments: list(type) >= 0")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
+ .Attr("preserve_cardinality: bool = false")
.SetShapeFn(shape_inference::ScalarShape);
REGISTER_OP("ExperimentalSetStatsAggregatorDataset")
@@ -425,6 +428,7 @@
.Attr("Targuments: list(type) >= 0")
.Attr("output_types: list(type) >= 1")
.Attr("output_shapes: list(shape) >= 1")
+ .Attr("preserve_cardinality: bool = false")
.SetShapeFn([](shape_inference::InferenceContext* c) {
// Use index from the end to retrieve the Input shapes,
// so that to avoid guessing the length of "other_arguments".
diff --git a/tensorflow/core/ops/linalg_ops.cc b/tensorflow/core/ops/linalg_ops.cc
index 525b19e..952ee4b 100644
--- a/tensorflow/core/ops/linalg_ops.cc
+++ b/tensorflow/core/ops/linalg_ops.cc
@@ -109,6 +109,30 @@
return Status::OK();
}
+// Input is [...,N,N].
+// First and second outputs are:
+// [...,N,N]; [...,N].
+Status LuShapeFn(InferenceContext* c) {
+ ShapeHandle input;
+ TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input));
+
+ DimensionHandle n;
+ TF_RETURN_IF_ERROR(c->Merge(c->Dim(input, -2), c->Dim(input, -1), &n));
+
+ ShapeHandle batch_shape;
+ TF_RETURN_IF_ERROR(c->Subshape(input, 0, -2, &batch_shape));
+
+ ShapeHandle lu_shape;
+ ShapeHandle p_shape;
+
+ TF_RETURN_IF_ERROR(c->Concatenate(batch_shape, c->Matrix(n, n), &lu_shape));
+ TF_RETURN_IF_ERROR(c->Concatenate(batch_shape, c->Vector(n), &p_shape));
+
+ c->set_output(0, lu_shape);
+ c->set_output(1, p_shape);
+ return Status::OK();
+}
+
// Input is [...,M,N].
// First and second outputs are:
// [...,M,M]; [...,M,N], if full_matrices is true,
@@ -289,6 +313,14 @@
.Attr("T: {double, float, complex64, complex128}")
.SetShapeFn(SelfAdjointEigV2ShapeFn);
+REGISTER_OP("Lu")
+ .Input("input: T")
+ .Output("lu: T")
+ .Output("p: output_idx_type")
+ .Attr("T: {double, float, complex64, complex128}")
+ .Attr("output_idx_type: {int32, int64} = DT_INT32")
+ .SetShapeFn(LuShapeFn);
+
REGISTER_OP("MatrixSolve")
.Input("matrix: T")
.Input("rhs: T")
diff --git a/tensorflow/core/ops/linalg_ops_test.cc b/tensorflow/core/ops/linalg_ops_test.cc
index f4be820..bfacee1 100644
--- a/tensorflow/core/ops/linalg_ops_test.cc
+++ b/tensorflow/core/ops/linalg_ops_test.cc
@@ -274,4 +274,23 @@
INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, "[1]");
}
+TEST(LinalgOpsTest, Lu_ShapeFn) {
+ ShapeInferenceTestOp op("Lu");
+ INFER_OK(op, "?", "?;?");
+ INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, "[1]");
+ INFER_ERROR("Dimensions must be equal, but are 1 and 2", op, "[1,?,3,4,1,2]");
+
+ INFER_OK(op, "[?,?]", "[d0_0,d0_0];[d0_0]");
+ INFER_OK(op, "[1,?]", "[d0_0,d0_0];[d0_0]");
+ INFER_OK(op, "[?,1]", "[d0_1,d0_1];[d0_1]");
+
+ // Repeat previous block of tests with input rank > 2.
+ INFER_OK(op, "[1,?,3,4,?,?]",
+ "[d0_0,d0_1,d0_2,d0_3,d0_4,d0_4];[d0_0,d0_1,d0_2,d0_3,d0_4]");
+ INFER_OK(op, "[1,?,3,4,1,?]",
+ "[d0_0,d0_1,d0_2,d0_3,d0_4,d0_4];[d0_0,d0_1,d0_2,d0_3,d0_4]");
+ INFER_OK(op, "[1,?,3,4,?,1]",
+ "[d0_0,d0_1,d0_2,d0_3,d0_5,d0_5];[d0_0,d0_1,d0_2,d0_3,d0_5]");
+}
+
} // end namespace tensorflow
diff --git a/tensorflow/core/ops/mkl_nn_ops.cc b/tensorflow/core/ops/mkl_nn_ops.cc
index 9be3470..658afd9 100644
--- a/tensorflow/core/ops/mkl_nn_ops.cc
+++ b/tensorflow/core/ops/mkl_nn_ops.cc
@@ -32,6 +32,33 @@
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;
+REGISTER_OP("_MklFusedConv2D")
+ .Input("input: T")
+ .Input("filter: T")
+ .Input("args: num_args * T")
+ .Input("mkl_input: uint8")
+ .Input("mkl_filter: uint8")
+ .Input("mkl_args: num_args * uint8")
+ .Output("output: T")
+ .Output("filter_output: T")
+ .Output("mkl_output: uint8")
+ .Output("mkl_filter_output: uint8")
+ .Attr("T: {float}")
+ .Attr("num_args: int >= 0")
+ .Attr("strides: list(int)")
+ .Attr(GetPaddingAttrString())
+ .Attr(GetConvnetDataFormatAttrString())
+ .Attr("dilations: list(int) = [1, 1, 1, 1]")
+ .Attr("fused_ops: list(string) = []")
+ // Attributes for the FusedBatchNorm ------------------------------------ //
+ .Attr("epsilon: float = 0.0001")
+ // ---------------------------------------------------------------------- //
+ .SetShapeFn(shape_inference::Conv2DShape)
+ .Doc(R"doc(
+*NOTE*: Do not invoke this operator directly in Python. MKL DNN graph transformer
+ is expected to create these operators.
+)doc");
+
REGISTER_OP("_MklQuantizedMaxPool")
.Input("input: T")
.Input("min_input: float")
diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc
index cb417ea..ee528c7 100644
--- a/tensorflow/core/ops/nn_ops.cc
+++ b/tensorflow/core/ops/nn_ops.cc
@@ -1609,6 +1609,55 @@
expected to invoke these operators.
)doc");
+REGISTER_OP("__MklDummyPadWithConv2D")
+ .Input("input: T")
+ .Input("filter: T")
+ .Input("paddings: Tpaddings")
+ .Output("output: T")
+ .Attr("T: {half, float, double}")
+ .Attr("strides: list(int)")
+ .Attr("use_cudnn_on_gpu: bool = true")
+ .Attr(GetPaddingAttrString())
+ .Attr(GetConvnetDataFormatAttrString())
+ .Attr("dilations: list(int) = [1, 1, 1, 1]")
+ .Attr("Tpaddings: {int32, int64} = DT_INT32")
+ .SetShapeFn(shape_inference::Conv2DShape)
+ .Doc(R"doc(
+Dummy node that enables fusing Pad and Conv2D operator for MKL. This node
+does not perform anything. It is just created as an intermediate output of
+merging Pad and Conv2D.
+
+NOTE Do not invoke this operator directly in Python. Graph rewrite pass is
+expected to invoke these operators.
+)doc");
+
+REGISTER_OP("_MklPadWithConv2D")
+ .Input("input: T")
+ .Input("filter: T")
+ .Input("paddings: Tpaddings")
+ .Input("mkl_input: uint8")
+ .Input("mkl_filter: uint8")
+ .Input("mkl_paddings: uint8")
+ .Output("output: T")
+ .Output("filter_output: T")
+ .Output("mkl_output: uint8")
+ .Output("mkl_filter_output: uint8")
+ .Attr("T: {half, float, double}")
+ .Attr("strides: list(int)")
+ .Attr("use_cudnn_on_gpu: bool = true")
+ .Attr(GetPaddingAttrString())
+ .Attr(GetConvnetDataFormatAttrString())
+ .Attr("dilations: list(int) = [1, 1, 1, 1]")
+ .Attr("Tpaddings: {int32, int64} = DT_INT32")
+ .SetShapeFn(shape_inference::Conv2DShape)
+ .Doc(R"doc(
+MKL version of Pad and Conv2D operator. Uses MKL DNN APIs to perform
+Pad and 2D convolution to the output of convolution.
+
+NOTE Do not invoke this operator directly in Python. Graph rewrite pass is
+expected to invoke these operators.
+)doc");
+
REGISTER_OP("_MklConv2DBackpropFilter")
.Input("input: T")
.Input("filter_sizes: int32")
@@ -1915,6 +1964,40 @@
expected to invoke these operators.
)doc");
+REGISTER_OP("_MklLeakyRelu")
+ .Input("features: T")
+ .Input("mkl_features: uint8")
+ .Output("activations: T")
+ .Output("mkl_activations: uint8")
+ .Attr("T: {half, float, double} = DT_FLOAT")
+ .Attr("alpha: float = 0.2")
+ .SetShapeFn(shape_inference::UnchangedShape)
+ .Doc(R"doc(
+MKL version of LeakyRelu operator. Uses MKL DNN APIs to implement
+LeakyRelu operator.
+
+NOTE Do not invoke this operator directly in Python. Graph rewrite pass is
+expected to invoke these operators.
+)doc");
+
+REGISTER_OP("_MklLeakyReluGrad")
+ .Input("gradients: T")
+ .Input("features: T")
+ .Input("mkl_gradients: uint8")
+ .Input("mkl_features: uint8")
+ .Output("backprops: T")
+ .Output("mkl_backprops: uint8")
+ .Attr("T: {half, float, double} = DT_FLOAT")
+ .Attr("alpha: float = 0.2")
+ .SetShapeFn(shape_inference::MergeBothInputsShapeFn)
+ .Doc(R"doc(
+MKL version of LeakyReluGrad operator. Uses MKL DNN APIs to compute rectified
+linear gradients for LeakyReluGrad operation.
+
+NOTE Do not invoke this operator directly in Python. Graph rewrite pass is
+expected to invoke these operators.
+)doc");
+
REGISTER_OP("_MklElu")
.Input("features: T")
.Input("mkl_features: uint8")
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index a9c712e..779d429 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -10091,6 +10091,17 @@
is_stateful: true
}
op {
+ name: "ExperimentalDatasetCardinality"
+ input_arg {
+ name: "input_dataset"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "cardinality"
+ type: DT_INT64
+ }
+}
+op {
name: "ExperimentalDatasetToTFRecord"
input_arg {
name: "input_dataset"
@@ -10493,6 +10504,13 @@
has_minimum: true
minimum: 1
}
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
}
op {
name: "ExperimentalMapDataset"
@@ -10536,6 +10554,13 @@
b: true
}
}
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
}
op {
name: "ExperimentalMatchingFilesDataset"
@@ -10674,6 +10699,13 @@
has_minimum: true
minimum: 1
}
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
}
op {
name: "ExperimentalParallelInterleaveDataset"
@@ -10908,6 +10940,13 @@
has_minimum: true
minimum: 1
}
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
}
op {
name: "ExperimentalSetStatsAggregatorDataset"
@@ -15752,6 +15791,46 @@
}
}
op {
+ name: "Lu"
+ input_arg {
+ name: "input"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "lu"
+ type_attr: "T"
+ }
+ output_arg {
+ name: "p"
+ type_attr: "output_idx_type"
+ }
+ attr {
+ name: "T"
+ type: "type"
+ allowed_values {
+ list {
+ type: DT_DOUBLE
+ type: DT_FLOAT
+ type: DT_COMPLEX64
+ type: DT_COMPLEX128
+ }
+ }
+ }
+ attr {
+ name: "output_idx_type"
+ type: "type"
+ default_value {
+ type: DT_INT32
+ }
+ allowed_values {
+ list {
+ type: DT_INT32
+ type: DT_INT64
+ }
+ }
+ }
+}
+op {
name: "MakeIterator"
input_arg {
name: "dataset"
@@ -15843,6 +15922,13 @@
b: true
}
}
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
}
op {
name: "MapDefun"
@@ -19893,6 +19979,13 @@
b: false
}
}
+ attr {
+ name: "preserve_cardinality"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
}
op {
name: "ParameterizedTruncatedNormal"
@@ -37314,6 +37407,53 @@
}
}
op {
+ name: "UnicodeDecode"
+ input_arg {
+ name: "input"
+ type: DT_STRING
+ }
+ output_arg {
+ name: "row_splits"
+ type: DT_INT64
+ }
+ output_arg {
+ name: "char_values"
+ type: DT_INT32
+ }
+ attr {
+ name: "input_encoding"
+ type: "string"
+ }
+ attr {
+ name: "errors"
+ type: "string"
+ default_value {
+ s: "replace"
+ }
+ allowed_values {
+ list {
+ s: "strict"
+ s: "replace"
+ s: "ignore"
+ }
+ }
+ }
+ attr {
+ name: "replacement_char"
+ type: "int"
+ default_value {
+ i: 65533
+ }
+ }
+ attr {
+ name: "replace_control_characters"
+ type: "bool"
+ default_value {
+ b: false
+ }
+ }
+}
+op {
name: "UnicodeDecodeWithOffsets"
input_arg {
name: "input"
@@ -38061,6 +38201,17 @@
is_stateful: true
}
op {
+ name: "UnwrapDatasetVariant"
+ input_arg {
+ name: "input_handle"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "output_handle"
+ type: DT_VARIANT
+ }
+}
+op {
name: "UpperBound"
input_arg {
name: "sorted_inputs"
@@ -38379,6 +38530,17 @@
}
}
op {
+ name: "WrapDatasetVariant"
+ input_arg {
+ name: "input_handle"
+ type: DT_VARIANT
+ }
+ output_arg {
+ name: "output_handle"
+ type: DT_VARIANT
+ }
+}
+op {
name: "WriteAudioSummary"
input_arg {
name: "writer"
diff --git a/tensorflow/core/ops/string_ops.cc b/tensorflow/core/ops/string_ops.cc
index 8ea74f1..d012ce6 100644
--- a/tensorflow/core/ops/string_ops.cc
+++ b/tensorflow/core/ops/string_ops.cc
@@ -296,6 +296,27 @@
.Attr("replace_control_characters: bool = false")
.SetShapeFn(shape_inference::UnchangedShape);
+REGISTER_OP("UnicodeDecode")
+ .Input("input: string")
+ .Output("row_splits: int64")
+ .Output("char_values: int32")
+ .Attr("input_encoding: string")
+ .Attr("errors: {'strict', 'replace', 'ignore'} = 'replace'")
+ .Attr("replacement_char: int = 65533") // 0xFFFD unicode replacement char
+ .Attr("replace_control_characters: bool = false")
+ .SetShapeFn([](InferenceContext* c) {
+ // row_splits.shape == [input.size() + 1]
+ DimensionHandle num_row_splits;
+ DimensionHandle input_size = c->NumElements(c->input(0));
+ TF_RETURN_IF_ERROR(c->Add(input_size, 1, &num_row_splits));
+ c->set_output(0, c->Vector(num_row_splits));
+
+ // char_values.shape == [num_chars]
+ DimensionHandle num_chars = c->UnknownDim();
+ c->set_output(1, c->Vector(num_chars));
+ return Status::OK();
+ });
+
REGISTER_OP("UnicodeDecodeWithOffsets")
.Input("input: string")
.Output("row_splits: int64")
diff --git a/tensorflow/core/platform/cpu_info.h b/tensorflow/core/platform/cpu_info.h
index 6eba832..c9208cc 100644
--- a/tensorflow/core/platform/cpu_info.h
+++ b/tensorflow/core/platform/cpu_info.h
@@ -32,9 +32,22 @@
// Returns an estimate of the number of schedulable CPUs for this
// process. Usually, it's constant throughout the lifetime of a
// process, but it might change if the underlying cluster management
-// software can change it dynamically.
+// software can change it dynamically. If the underlying call fails, a default
+// value (e.g. `4`) may be returned.
int NumSchedulableCPUs();
+// Returns the total number of CPUs on the system. This number should
+// not change even if the underlying cluster management software may
+// change the number of schedulable CPUs. Unlike `NumSchedulableCPUs`, if the
+// underlying call fails, an invalid value of -1 will be returned;
+// the user must check for validity.
+static constexpr int kUnknownCPU = -1;
+int NumTotalCPUs();
+
+// Returns the id of the current CPU. Returns -1 if the current CPU cannot be
+// identified. If successful, the return value will be in [0, NumTotalCPUs()).
+int GetCurrentCPU();
+
// Returns an estimate of the number of hyperthreads per physical core
// on the CPU
int NumHyperthreadsPerCore();
diff --git a/tensorflow/core/platform/cuda_libdevice_path.cc b/tensorflow/core/platform/cuda_libdevice_path.cc
deleted file mode 100644
index 4d6532b..0000000
--- a/tensorflow/core/platform/cuda_libdevice_path.cc
+++ /dev/null
@@ -1,26 +0,0 @@
-/* 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/core/platform/cuda_libdevice_path.h"
-
-#include "tensorflow/core/lib/io/path.h"
-
-namespace tensorflow {
-
-string LibdeviceRoot() {
- return tensorflow::io::JoinPath(tensorflow::CudaRoot(), "nvvm/libdevice");
-}
-
-} // namespace tensorflow
diff --git a/tensorflow/core/platform/cuda_libdevice_path.h b/tensorflow/core/platform/cuda_libdevice_path.h
index 6ef565e..f2dbff9 100644
--- a/tensorflow/core/platform/cuda_libdevice_path.h
+++ b/tensorflow/core/platform/cuda_libdevice_path.h
@@ -16,16 +16,14 @@
#ifndef TENSORFLOW_CORE_PLATFORM_CUDA_LIBDEVICE_PATH_H_
#define TENSORFLOW_CORE_PLATFORM_CUDA_LIBDEVICE_PATH_H_
+#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
-// Returns the root directory of the CUDA SDK, which contains sub-folders such
-// as bin, lib64, and nvvm.
-string CudaRoot();
-
-// Returns the directory that contains nvvm libdevice files in the CUDA SDK.
-string LibdeviceRoot();
+// Returns, in order of preference, potential locations of the root directory of
+// the CUDA SDK, which contains sub-folders such as bin, lib64, and nvvm.
+std::vector<string> CandidateCudaRoots();
} // namespace tensorflow
diff --git a/tensorflow/core/platform/cuda_libdevice_path_test.cc b/tensorflow/core/platform/cuda_libdevice_path_test.cc
deleted file mode 100644
index 2d34239..0000000
--- a/tensorflow/core/platform/cuda_libdevice_path_test.cc
+++ /dev/null
@@ -1,35 +0,0 @@
-/* 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/core/platform/cuda_libdevice_path.h"
-
-#include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/lib/io/path.h"
-#include "tensorflow/core/platform/env.h"
-#include "tensorflow/core/platform/test.h"
-
-namespace tensorflow {
-
-#if GOOGLE_CUDA
-TEST(CudaLibdevicePathTest, LibdevicePath) {
- VLOG(2) << "Libdevice root = " << LibdeviceRoot();
- std::vector<string> libdevice_files;
- TF_EXPECT_OK(Env::Default()->GetMatchingPaths(
- io::JoinPath(LibdeviceRoot(), "libdevice.*.bc"), &libdevice_files));
- EXPECT_LT(0, libdevice_files.size());
-}
-#endif
-
-} // namespace tensorflow
diff --git a/tensorflow/core/platform/default/cuda_libdevice_path.cc b/tensorflow/core/platform/default/cuda_libdevice_path.cc
index 20ee3ad..a8b2e72 100644
--- a/tensorflow/core/platform/default/cuda_libdevice_path.cc
+++ b/tensorflow/core/platform/default/cuda_libdevice_path.cc
@@ -16,6 +16,7 @@
#include "tensorflow/core/platform/cuda_libdevice_path.h"
#include <stdlib.h>
+#include <vector>
#if !defined(PLATFORM_GOOGLE)
#include "cuda/cuda_config.h"
@@ -24,9 +25,9 @@
namespace tensorflow {
-string CudaRoot() {
+std::vector<string> CandidateCudaRoots() {
VLOG(3) << "CUDA root = " << TF_CUDA_TOOLKIT_PATH;
- return TF_CUDA_TOOLKIT_PATH;
+ return {TF_CUDA_TOOLKIT_PATH};
}
} // namespace tensorflow
diff --git a/tensorflow/core/platform/default/human_readable_json.cc b/tensorflow/core/platform/default/human_readable_json.cc
index 9f97c82..bf9c7b7 100644
--- a/tensorflow/core/platform/default/human_readable_json.cc
+++ b/tensorflow/core/platform/default/human_readable_json.cc
@@ -20,7 +20,7 @@
namespace tensorflow {
-Status ProtoToHumanReadableJson(const ::google::protobuf::Message& proto,
+Status ProtoToHumanReadableJson(const protobuf::Message& proto,
string* result) {
#ifdef TENSORFLOW_LITE_PROTOS
*result = "[human readable output not available on Android]";
@@ -28,7 +28,7 @@
#else
result->clear();
- auto status = google::protobuf::util::MessageToJsonString(proto, result);
+ auto status = protobuf::util::MessageToJsonString(proto, result);
if (!status.ok()) {
// Convert error_msg google::protobuf::StringPiece to
// tensorflow::StringPiece.
@@ -41,8 +41,7 @@
#endif
}
-Status HumanReadableJsonToProto(const string& str,
- ::google::protobuf::Message* proto) {
+Status HumanReadableJsonToProto(const string& str, protobuf::Message* proto) {
#ifdef TENSORFLOW_LITE_PROTOS
return errors::Internal("Cannot parse JSON protos on Android");
#else
diff --git a/tensorflow/core/platform/env.h b/tensorflow/core/platform/env.h
index 7374fcc..1b53828 100644
--- a/tensorflow/core/platform/env.h
+++ b/tensorflow/core/platform/env.h
@@ -167,11 +167,24 @@
Status DeleteFile(const string& fname);
/// \brief Deletes the specified directory and all subdirectories and files
- /// underneath it. undeleted_files and undeleted_dirs stores the number of
- /// files and directories that weren't deleted (unspecified if the return
- /// status is not OK).
+ /// underneath it. This is accomplished by traversing the directory tree
+ /// rooted at dirname and deleting entries as they are encountered.
+ ///
+ /// If dirname itself is not readable or does not exist, *undeleted_dir_count
+ /// is set to 1, *undeleted_file_count is set to 0 and an appropriate status
+ /// (e.g. NOT_FOUND) is returned.
+ ///
+ /// If dirname and all its descendants were successfully deleted, TF_OK is
+ /// returned and both error counters are set to zero.
+ ///
+ /// Otherwise, while traversing the tree, undeleted_file_count and
+ /// undeleted_dir_count are updated if an entry of the corresponding type
+ /// could not be deleted. The returned error status represents the reason that
+ /// any one of these entries could not be deleted.
+ ///
/// REQUIRES: undeleted_files, undeleted_dirs to be not null.
- /// Typical return codes
+ ///
+ /// Typical return codes:
/// * OK - dirname exists and we were able to delete everything underneath.
/// * NOT_FOUND - dirname doesn't exist
/// * PERMISSION_DENIED - dirname or some descendant is not writable
diff --git a/tensorflow/core/platform/file_system.h b/tensorflow/core/platform/file_system.h
index 156af6c..c84a93b 100644
--- a/tensorflow/core/platform/file_system.h
+++ b/tensorflow/core/platform/file_system.h
@@ -167,10 +167,23 @@
virtual Status DeleteDir(const string& dirname) = 0;
/// \brief Deletes the specified directory and all subdirectories and files
- /// underneath it. undeleted_files and undeleted_dirs stores the number of
- /// files and directories that weren't deleted (unspecified if the return
- /// status is not OK).
+ /// underneath it. This is accomplished by traversing the directory tree
+ /// rooted at dirname and deleting entries as they are encountered.
+ ///
+ /// If dirname itself is not readable or does not exist, *undeleted_dir_count
+ /// is set to 1, *undeleted_file_count is set to 0 and an appropriate status
+ /// (e.g. NOT_FOUND) is returned.
+ ///
+ /// If dirname and all its descendants were successfully deleted, TF_OK is
+ /// returned and both error counters are set to zero.
+ ///
+ /// Otherwise, while traversing the tree, undeleted_file_count and
+ /// undeleted_dir_count are updated if an entry of the corresponding type
+ /// could not be deleted. The returned error status represents the reason that
+ /// any one of these entries could not be deleted.
+ ///
/// REQUIRES: undeleted_files, undeleted_dirs to be not null.
+ ///
/// Typical return codes:
/// * OK - dirname exists and we were able to delete everything underneath.
/// * NOT_FOUND - dirname doesn't exist
diff --git a/tensorflow/core/platform/port_test.cc b/tensorflow/core/platform/port_test.cc
index 15c3cb2..9d144ef 100644
--- a/tensorflow/core/platform/port_test.cc
+++ b/tensorflow/core/platform/port_test.cc
@@ -33,6 +33,12 @@
}
}
+TEST(Port, GetCurrentCPU) {
+ const int cpu = GetCurrentCPU();
+ EXPECT_GE(cpu, 0);
+ EXPECT_LT(cpu, NumTotalCPUs());
+}
+
TEST(ConditionVariable, WaitForMilliseconds_Timeout) {
mutex m;
mutex_lock l(m);
@@ -78,3 +84,9 @@
} // namespace port
} // namespace tensorflow
+
+int main(int argc, char** argv) {
+ // On Linux, add: FLAGS_logtostderr = true;
+ ::testing::InitGoogleTest(&argc, argv);
+ return RUN_ALL_TESTS();
+}
diff --git a/tensorflow/core/platform/posix/port.cc b/tensorflow/core/platform/posix/port.cc
index acdd779..0fac8b1 100644
--- a/tensorflow/core/platform/posix/port.cc
+++ b/tensorflow/core/platform/posix/port.cc
@@ -25,7 +25,14 @@
#if defined(__linux__) && !defined(__ANDROID__)
#include <sched.h>
#include <sys/sysinfo.h>
+#else
+#include <sys/syscall.h>
#endif
+
+#if !defined(__APPLE__) && (__x86_64__ || __i386__)
+#include <cpuid.h>
+#endif
+
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
@@ -69,6 +76,34 @@
return kDefaultCores;
}
+int NumTotalCPUs() {
+ int count = absl::base_internal::NumCPUs();
+ return (count == 0) ? kUnknownCPU : count;
+}
+
+int GetCurrentCPU() {
+#if defined(__linux__) && !defined(__ANDROID__)
+ return sched_getcpu();
+#elif defined(__cpuid_count)
+ // Attempt to use cpuid on all other platforms. If that fails, perform a
+ // syscall.
+ uint32_t eax, ebx, ecx, edx;
+ __cpuid_count(/*leaf=*/1, /*subleaf=*/0, eax, ebx, ecx, edx);
+ if ((edx & (1 << 9)) != 0) {
+ // EBX bits 24-31 are APIC ID
+ return static_cast<unsigned int>(ebx >> 24);
+ }
+#elif defined(__NR_getcpu)
+ unsigned int cpu;
+ if (syscall(__NR_getcpu, &cpu, NULL, NULL) < 0) {
+ return kUnknownCPU;
+ } else {
+ return static_cast<int>(cpu);
+ }
+#endif
+ return kUnknownCPU;
+}
+
int NumHyperthreadsPerCore() {
static const int ht_per_core = tensorflow::port::CPUIDNumSMT();
return (ht_per_core > 0) ? ht_per_core : 1;
@@ -83,9 +118,7 @@
void NUMASetThreadNodeAffinity(int node) {}
-int NUMAGetThreadNodeAffinity() {
- return kNUMANoAffinity;
-}
+int NUMAGetThreadNodeAffinity() { return kNUMANoAffinity; }
void* AlignedMalloc(size_t size, int minimum_alignment) {
#if defined(__ANDROID__)
diff --git a/tensorflow/core/platform/windows/port.cc b/tensorflow/core/platform/windows/port.cc
index 911ea19..b902c85 100644
--- a/tensorflow/core/platform/windows/port.cc
+++ b/tensorflow/core/platform/windows/port.cc
@@ -21,6 +21,7 @@
#endif
#include <Windows.h>
+#include <processthreadsapi.h>
#include <shlwapi.h>
#include "tensorflow/core/platform/cpu_info.h"
@@ -54,6 +55,30 @@
return system_info.dwNumberOfProcessors;
}
+int NumTotalCPUs() {
+ // TODO(ebrevdo): Make this more accurate.
+ //
+ // This only returns the number of processors in the current
+ // processor group; which may be undercounting if you have more than 64 cores.
+ // For that case, one needs to call
+ // GetLogicalProcessorInformationEx(RelationProcessorCore, ...) and accumulate
+ // the Size fields by iterating over the written-to buffer. Since I can't
+ // easily test this on Windows, I'm deferring this to someone who can!
+ //
+ // If you fix this, also consider updatig GetCurrentCPU below.
+ return NumSchedulableCPUs();
+}
+
+int GetCurrentCPU() {
+ // NOTE(ebrevdo): This returns the processor number within the processor
+ // group on systems with >64 processors. Therefore it doesn't necessarily map
+ // naturally to an index in NumSchedulableCPUs().
+ //
+ // On the plus side, this number is probably guaranteed to be within
+ // [0, NumTotalCPUs()) due to its incomplete implementation.
+ return GetCurrentProcessorNumber();
+}
+
bool NUMAEnabled() {
// Not yet implemented: coming soon.
return false;
diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 07eeeb4..a55fe17 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -24,7 +24,7 @@
// TF_VERSION_SUFFIX is non-empty for pre-releases (e.g. "-alpha", "-alpha.1",
// "-beta", "-rc", "-rc.1")
-#define TF_VERSION_SUFFIX "-rc0"
+#define TF_VERSION_SUFFIX ""
#define TF_STR_HELPER(x) #x
#define TF_STR(x) TF_STR_HELPER(x)
diff --git a/tensorflow/core/util/cuda_launch_config.h b/tensorflow/core/util/cuda_launch_config.h
index 080d406..c0ae634 100644
--- a/tensorflow/core/util/cuda_launch_config.h
+++ b/tensorflow/core/util/cuda_launch_config.h
@@ -21,7 +21,6 @@
#include <algorithm>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
-#include "cuda/include/cuda.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/stream_executor.h"
diff --git a/tensorflow/core/util/port.cc b/tensorflow/core/util/port.cc
index e01058d..7dc8ddd 100644
--- a/tensorflow/core/util/port.cc
+++ b/tensorflow/core/util/port.cc
@@ -15,9 +15,6 @@
#include "tensorflow/core/util/port.h"
-#if GOOGLE_CUDA
-#include "cuda/include/cuda.h"
-#endif
namespace tensorflow {
diff --git a/tensorflow/examples/android/README.md b/tensorflow/examples/android/README.md
index 82bc3ff..4e4e168 100644
--- a/tensorflow/examples/android/README.md
+++ b/tensorflow/examples/android/README.md
@@ -180,7 +180,7 @@
build the APK. Run this from your workspace root:
```bash
-bazel build -c opt //tensorflow/examples/android:tensorflow_demo
+bazel build --cxxopt='--std=c++11' -c opt //tensorflow/examples/android:tensorflow_demo
```
##### Install
diff --git a/tensorflow/examples/tutorials/mnist/mnist_with_summaries.py b/tensorflow/examples/tutorials/mnist/mnist_with_summaries.py
index 7967e22..1854e84 100644
--- a/tensorflow/examples/tutorials/mnist/mnist_with_summaries.py
+++ b/tensorflow/examples/tutorials/mnist/mnist_with_summaries.py
@@ -183,7 +183,8 @@
if tf.gfile.Exists(FLAGS.log_dir):
tf.gfile.DeleteRecursively(FLAGS.log_dir)
tf.gfile.MakeDirs(FLAGS.log_dir)
- train()
+ with tf.Graph().as_default():
+ train()
if __name__ == '__main__':
diff --git a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py
index b09ee99..77889ef 100644
--- a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py
+++ b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py
@@ -91,7 +91,7 @@
def build_dataset(words, n_words):
"""Process raw inputs into a dataset."""
- count = [['UNK', -1]]
+ count = [('UNK', -1)]
count.extend(collections.Counter(words).most_common(n_words - 1))
dictionary = dict()
for word, _ in count:
diff --git a/tensorflow/examples/wav_to_spectrogram/wav_to_spectrogram.cc b/tensorflow/examples/wav_to_spectrogram/wav_to_spectrogram.cc
index 4a42983..464484d 100644
--- a/tensorflow/examples/wav_to_spectrogram/wav_to_spectrogram.cc
+++ b/tensorflow/examples/wav_to_spectrogram/wav_to_spectrogram.cc
@@ -33,7 +33,6 @@
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
using tensorflow::DT_FLOAT;
using tensorflow::DT_UINT8;
diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 5cd0629..4624d12 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -5588,6 +5588,467 @@
return op.Output(0)
}
+// MapSizeAttr is an optional argument to MapSize.
+type MapSizeAttr func(optionalAttr)
+
+// MapSizeCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapSizeCapacity(value int64) MapSizeAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// MapSizeMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapSizeMemoryLimit(value int64) MapSizeAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// MapSizeContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func MapSizeContainer(value string) MapSizeAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// MapSizeSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func MapSizeSharedName(value string) MapSizeAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op returns the number of elements in the underlying container.
+func MapSize(scope *Scope, dtypes []tf.DataType, optional ...MapSizeAttr) (size tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MapSize",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// MapUnstageAttr is an optional argument to MapUnstage.
+type MapUnstageAttr func(optionalAttr)
+
+// MapUnstageCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapUnstageCapacity(value int64) MapUnstageAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// MapUnstageMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapUnstageMemoryLimit(value int64) MapUnstageAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// MapUnstageContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func MapUnstageContainer(value string) MapUnstageAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// MapUnstageSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func MapUnstageSharedName(value string) MapUnstageAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op removes and returns the values associated with the key
+//
+// from the underlying container. If the underlying container
+// does not contain this key, the op will block until it does.
+func MapUnstage(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.DataType, optional ...MapUnstageAttr) (values []tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MapUnstage",
+ Input: []tf.Input{
+ key, indices,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ if scope.Err() != nil {
+ return
+ }
+ var idx int
+ var err error
+ if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
+ scope.UpdateErr("MapUnstage", err)
+ return
+ }
+ return values
+}
+
+// MapPeekAttr is an optional argument to MapPeek.
+type MapPeekAttr func(optionalAttr)
+
+// MapPeekCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapPeekCapacity(value int64) MapPeekAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// MapPeekMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapPeekMemoryLimit(value int64) MapPeekAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// MapPeekContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func MapPeekContainer(value string) MapPeekAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// MapPeekSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func MapPeekSharedName(value string) MapPeekAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op peeks at the values at the specified key. If the
+//
+// underlying container does not contain this key
+// this op will block until it does.
+func MapPeek(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.DataType, optional ...MapPeekAttr) (values []tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MapPeek",
+ Input: []tf.Input{
+ key, indices,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ if scope.Err() != nil {
+ return
+ }
+ var idx int
+ var err error
+ if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
+ scope.UpdateErr("MapPeek", err)
+ return
+ }
+ return values
+}
+
+// MapStageAttr is an optional argument to MapStage.
+type MapStageAttr func(optionalAttr)
+
+// MapStageCapacity sets the optional capacity attribute to value.
+//
+// value: Maximum number of elements in the Staging Area. If > 0, inserts
+// on the container will block when the capacity is reached.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapStageCapacity(value int64) MapStageAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// MapStageMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapStageMemoryLimit(value int64) MapStageAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// MapStageContainer sets the optional container attribute to value.
+//
+// value: If non-empty, this queue is placed in the given container. Otherwise,
+// a default container is used.
+// If not specified, defaults to ""
+func MapStageContainer(value string) MapStageAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// MapStageSharedName sets the optional shared_name attribute to value.
+//
+// value: It is necessary to match this name to the matching Unstage Op.
+// If not specified, defaults to ""
+func MapStageSharedName(value string) MapStageAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Stage (key, values) in the underlying container which behaves like a hashtable.
+//
+// Arguments:
+// key: int64
+//
+// values: a list of tensors
+// dtypes A list of data types that inserted values should adhere to.
+//
+//
+// Returns the created operation.
+func MapStage(scope *Scope, key tf.Output, indices tf.Output, values []tf.Output, dtypes []tf.DataType, optional ...MapStageAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MapStage",
+ Input: []tf.Input{
+ key, indices, tf.OutputList(values),
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// StageClearAttr is an optional argument to StageClear.
+type StageClearAttr func(optionalAttr)
+
+// StageClearCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func StageClearCapacity(value int64) StageClearAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// StageClearMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func StageClearMemoryLimit(value int64) StageClearAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// StageClearContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func StageClearContainer(value string) StageClearAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// StageClearSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func StageClearSharedName(value string) StageClearAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op removes all elements in the underlying container.
+//
+// Returns the created operation.
+func StageClear(scope *Scope, dtypes []tf.DataType, optional ...StageClearAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "StageClear",
+
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// StageSizeAttr is an optional argument to StageSize.
+type StageSizeAttr func(optionalAttr)
+
+// StageSizeCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func StageSizeCapacity(value int64) StageSizeAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// StageSizeMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func StageSizeMemoryLimit(value int64) StageSizeAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// StageSizeContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func StageSizeContainer(value string) StageSizeAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// StageSizeSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func StageSizeSharedName(value string) StageSizeAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op returns the number of elements in the underlying container.
+func StageSize(scope *Scope, dtypes []tf.DataType, optional ...StageSizeAttr) (size tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "StageSize",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// StagePeekAttr is an optional argument to StagePeek.
+type StagePeekAttr func(optionalAttr)
+
+// StagePeekCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func StagePeekCapacity(value int64) StagePeekAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// StagePeekMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func StagePeekMemoryLimit(value int64) StagePeekAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// StagePeekContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func StagePeekContainer(value string) StagePeekAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// StagePeekSharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func StagePeekSharedName(value string) StagePeekAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op peeks at the values at the specified index. If the
+//
+// underlying container does not contain sufficient elements
+// this op will block until it does. This Op is optimized for
+// performance.
+func StagePeek(scope *Scope, index tf.Output, dtypes []tf.DataType, optional ...StagePeekAttr) (values []tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "StagePeek",
+ Input: []tf.Input{
+ index,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ if scope.Err() != nil {
+ return
+ }
+ var idx int
+ var err error
+ if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
+ scope.UpdateErr("StagePeek", err)
+ return
+ }
+ return values
+}
+
// Compute the regularized incomplete beta integral \\(I_x(a, b)\\).
//
// The regularized incomplete beta integral is defined as:
@@ -5824,37 +6285,6 @@
return op.Output(0)
}
-// ApproximateEqualAttr is an optional argument to ApproximateEqual.
-type ApproximateEqualAttr func(optionalAttr)
-
-// ApproximateEqualTolerance sets the optional tolerance attribute to value.
-// If not specified, defaults to 1e-05
-func ApproximateEqualTolerance(value float32) ApproximateEqualAttr {
- return func(m optionalAttr) {
- m["tolerance"] = value
- }
-}
-
-// Returns the truth value of abs(x-y) < tolerance element-wise.
-func ApproximateEqual(scope *Scope, x tf.Output, y tf.Output, optional ...ApproximateEqualAttr) (z tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ApproximateEqual",
- Input: []tf.Input{
- x, y,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Returns x / y element-wise.
//
// *NOTE*: `Div` supports broadcasting. More about broadcasting
@@ -7654,44 +8084,6 @@
return op.Output(0)
}
-// Computes softsign gradients for a softsign operation.
-//
-// Arguments:
-// gradients: The backpropagated gradients to the corresponding softsign operation.
-// features: The features passed as input to the corresponding softsign operation.
-//
-// Returns The gradients: `gradients / (1 + abs(features)) ** 2`.
-func SoftsignGrad(scope *Scope, gradients tf.Output, features tf.Output) (backprops tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SoftsignGrad",
- Input: []tf.Input{
- gradients, features,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Provides the time since epoch in seconds.
-//
-// Returns the timestamp as a `float64` for seconds since the Unix epoch.
-//
-// Note: the timestamp is computed when the op is executed, not when it is added
-// to the graph.
-func Timestamp(scope *Scope) (ts tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Timestamp",
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// BatchMatMulAttr is an optional argument to BatchMatMul.
type BatchMatMulAttr func(optionalAttr)
@@ -11229,24 +11621,6 @@
return op.Output(0)
}
-// Returns the truth value of (x > y) element-wise.
-//
-// *NOTE*: `Greater` supports broadcasting. More about broadcasting
-// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
-func Greater(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Greater",
- Input: []tf.Input{
- x, y,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// ResourceSparseApplyRMSPropAttr is an optional argument to ResourceSparseApplyRMSProp.
type ResourceSparseApplyRMSPropAttr func(optionalAttr)
@@ -11950,6 +12324,175 @@
return scope.AddOperation(opspec)
}
+// MaxPoolAttr is an optional argument to MaxPool.
+type MaxPoolAttr func(optionalAttr)
+
+// MaxPoolDataFormat sets the optional data_format attribute to value.
+//
+// value: Specify the data format of the input and output data. With the
+// default format "NHWC", the data is stored in the order of:
+// [batch, in_height, in_width, in_channels].
+// Alternatively, the format could be "NCHW", the data storage order of:
+// [batch, in_channels, in_height, in_width].
+// If not specified, defaults to "NHWC"
+func MaxPoolDataFormat(value string) MaxPoolAttr {
+ return func(m optionalAttr) {
+ m["data_format"] = value
+ }
+}
+
+// Performs max pooling on the input.
+//
+// Arguments:
+// input: 4-D input to pool over.
+// ksize: The size of the window for each dimension of the input tensor.
+// strides: The stride of the sliding window for each dimension of the
+// input tensor.
+// padding: The type of padding algorithm to use.
+//
+// Returns The max pooled output tensor.
+func MaxPool(scope *Scope, input tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPoolAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MaxPool",
+ Input: []tf.Input{
+ input,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Multiplies sparse updates into the variable referenced by `resource`.
+//
+// This operation computes
+//
+// # Scalar indices
+// ref[indices, ...] *= updates[...]
+//
+// # Vector indices (for each i)
+// ref[indices[i], ...] *= updates[i, ...]
+//
+// # High rank indices (for each i, ..., j)
+// ref[indices[i, ..., j], ...] *= updates[i, ..., j, ...]
+//
+// Duplicate entries are handled correctly: if multiple `indices` reference
+// the same location, their contributions multiply.
+//
+// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`.
+//
+// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
+// <img style="width:100%" src='https://www.tensorflow.org/images/ScatterAdd.png' alt>
+// </div>
+//
+// Arguments:
+// resource: Should be from a `Variable` node.
+// indices: A tensor of indices into the first dimension of `ref`.
+// updates: A tensor of updated values to add to `ref`.
+//
+// Returns the created operation.
+func ResourceScatterMul(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceScatterMul",
+ Input: []tf.Input{
+ resource, indices, updates,
+ },
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Subtracts sparse updates from the variable referenced by `resource`.
+//
+// This operation computes
+//
+// # Scalar indices
+// ref[indices, ...] -= updates[...]
+//
+// # Vector indices (for each i)
+// ref[indices[i], ...] -= updates[i, ...]
+//
+// # High rank indices (for each i, ..., j)
+// ref[indices[i, ..., j], ...] -= updates[i, ..., j, ...]
+//
+// Duplicate entries are handled correctly: if multiple `indices` reference
+// the same location, their contributions add.
+//
+// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`.
+//
+// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
+// <img style="width:100%" src='https://www.tensorflow.org/images/ScatterAdd.png' alt>
+// </div>
+//
+// Arguments:
+// resource: Should be from a `Variable` node.
+// indices: A tensor of indices into the first dimension of `ref`.
+// updates: A tensor of updated values to add to `ref`.
+//
+// Returns the created operation.
+func ResourceScatterSub(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceScatterSub",
+ Input: []tf.Input{
+ resource, indices, updates,
+ },
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Adds sparse updates to the variable referenced by `resource`.
+//
+// This operation computes
+//
+// # Scalar indices
+// ref[indices, ...] += updates[...]
+//
+// # Vector indices (for each i)
+// ref[indices[i], ...] += updates[i, ...]
+//
+// # High rank indices (for each i, ..., j)
+// ref[indices[i, ..., j], ...] += updates[i, ..., j, ...]
+//
+// Duplicate entries are handled correctly: if multiple `indices` reference
+// the same location, their contributions add.
+//
+// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`.
+//
+// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
+// <img style="width:100%" src='https://www.tensorflow.org/images/ScatterAdd.png' alt>
+// </div>
+//
+// Arguments:
+// resource: Should be from a `Variable` node.
+// indices: A tensor of indices into the first dimension of `ref`.
+// updates: A tensor of updated values to add to `ref`.
+//
+// Returns the created operation.
+func ResourceScatterAdd(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceScatterAdd",
+ Input: []tf.Input{
+ resource, indices, updates,
+ },
+ }
+ return scope.AddOperation(opspec)
+}
+
// ResourceSparseApplyFtrlV2Attr is an optional argument to ResourceSparseApplyFtrlV2.
type ResourceSparseApplyFtrlV2Attr func(optionalAttr)
@@ -12554,867 +13097,6 @@
return op.Output(0)
}
-// EncodeJpegAttr is an optional argument to EncodeJpeg.
-type EncodeJpegAttr func(optionalAttr)
-
-// EncodeJpegFormat sets the optional format attribute to value.
-//
-// value: Per pixel image format.
-// If not specified, defaults to ""
-func EncodeJpegFormat(value string) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["format"] = value
- }
-}
-
-// EncodeJpegQuality sets the optional quality attribute to value.
-//
-// value: Quality of the compression from 0 to 100 (higher is better and slower).
-// If not specified, defaults to 95
-func EncodeJpegQuality(value int64) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["quality"] = value
- }
-}
-
-// EncodeJpegProgressive sets the optional progressive attribute to value.
-//
-// value: If True, create a JPEG that loads progressively (coarse to fine).
-// If not specified, defaults to false
-func EncodeJpegProgressive(value bool) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["progressive"] = value
- }
-}
-
-// EncodeJpegOptimizeSize sets the optional optimize_size attribute to value.
-//
-// value: If True, spend CPU/RAM to reduce size with no quality change.
-// If not specified, defaults to false
-func EncodeJpegOptimizeSize(value bool) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["optimize_size"] = value
- }
-}
-
-// EncodeJpegChromaDownsampling sets the optional chroma_downsampling attribute to value.
-//
-// value: See http://en.wikipedia.org/wiki/Chroma_subsampling.
-// If not specified, defaults to true
-func EncodeJpegChromaDownsampling(value bool) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["chroma_downsampling"] = value
- }
-}
-
-// EncodeJpegDensityUnit sets the optional density_unit attribute to value.
-//
-// value: Unit used to specify `x_density` and `y_density`:
-// pixels per inch (`'in'`) or centimeter (`'cm'`).
-// If not specified, defaults to "in"
-func EncodeJpegDensityUnit(value string) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["density_unit"] = value
- }
-}
-
-// EncodeJpegXDensity sets the optional x_density attribute to value.
-//
-// value: Horizontal pixels per density unit.
-// If not specified, defaults to 300
-func EncodeJpegXDensity(value int64) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["x_density"] = value
- }
-}
-
-// EncodeJpegYDensity sets the optional y_density attribute to value.
-//
-// value: Vertical pixels per density unit.
-// If not specified, defaults to 300
-func EncodeJpegYDensity(value int64) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["y_density"] = value
- }
-}
-
-// EncodeJpegXmpMetadata sets the optional xmp_metadata attribute to value.
-//
-// value: If not empty, embed this XMP metadata in the image header.
-// If not specified, defaults to ""
-func EncodeJpegXmpMetadata(value string) EncodeJpegAttr {
- return func(m optionalAttr) {
- m["xmp_metadata"] = value
- }
-}
-
-// JPEG-encode an image.
-//
-// `image` is a 3-D uint8 Tensor of shape `[height, width, channels]`.
-//
-// The attr `format` can be used to override the color format of the encoded
-// output. Values can be:
-//
-// * `''`: Use a default format based on the number of channels in the image.
-// * `grayscale`: Output a grayscale JPEG image. The `channels` dimension
-// of `image` must be 1.
-// * `rgb`: Output an RGB JPEG image. The `channels` dimension
-// of `image` must be 3.
-//
-// If `format` is not specified or is the empty string, a default format is picked
-// in function of the number of channels in `image`:
-//
-// * 1: Output a grayscale image.
-// * 3: Output an RGB image.
-//
-// Arguments:
-// image: 3-D with shape `[height, width, channels]`.
-//
-// Returns 0-D. JPEG-encoded image.
-func EncodeJpeg(scope *Scope, image tf.Output, optional ...EncodeJpegAttr) (contents tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "EncodeJpeg",
- Input: []tf.Input{
- image,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// MultinomialAttr is an optional argument to Multinomial.
-type MultinomialAttr func(optionalAttr)
-
-// MultinomialSeed sets the optional seed attribute to value.
-//
-// value: If either seed or seed2 is set to be non-zero, the internal random number
-// generator is seeded by the given seed. Otherwise, a random seed is used.
-// If not specified, defaults to 0
-func MultinomialSeed(value int64) MultinomialAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// MultinomialSeed2 sets the optional seed2 attribute to value.
-//
-// value: A second seed to avoid seed collision.
-// If not specified, defaults to 0
-func MultinomialSeed2(value int64) MultinomialAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// MultinomialOutputDtype sets the optional output_dtype attribute to value.
-// If not specified, defaults to DT_INT64
-func MultinomialOutputDtype(value tf.DataType) MultinomialAttr {
- return func(m optionalAttr) {
- m["output_dtype"] = value
- }
-}
-
-// Draws samples from a multinomial distribution.
-//
-// Arguments:
-// logits: 2-D Tensor with shape `[batch_size, num_classes]`. Each slice `[i, :]`
-// represents the unnormalized log probabilities for all classes.
-// num_samples: 0-D. Number of independent samples to draw for each row slice.
-//
-// Returns 2-D Tensor with shape `[batch_size, num_samples]`. Each slice `[i, :]`
-// contains the drawn class labels with range `[0, num_classes)`.
-func Multinomial(scope *Scope, logits tf.Output, num_samples tf.Output, optional ...MultinomialAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "Multinomial",
- Input: []tf.Input{
- logits, num_samples,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// ResourceSparseApplyAdagradDAAttr is an optional argument to ResourceSparseApplyAdagradDA.
-type ResourceSparseApplyAdagradDAAttr func(optionalAttr)
-
-// ResourceSparseApplyAdagradDAUseLocking sets the optional use_locking attribute to value.
-//
-// value: If True, updating of the var and accum tensors will be protected by
-// a lock; otherwise the behavior is undefined, but may exhibit less contention.
-// If not specified, defaults to false
-func ResourceSparseApplyAdagradDAUseLocking(value bool) ResourceSparseApplyAdagradDAAttr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// Update entries in '*var' and '*accum' according to the proximal adagrad scheme.
-//
-// Arguments:
-// var_: Should be from a Variable().
-// gradient_accumulator: Should be from a Variable().
-// gradient_squared_accumulator: Should be from a Variable().
-// grad: The gradient.
-// indices: A vector of indices into the first dimension of var and accum.
-// lr: Learning rate. Must be a scalar.
-// l1: L1 regularization. Must be a scalar.
-// l2: L2 regularization. Must be a scalar.
-// global_step: Training step number. Must be a scalar.
-//
-// Returns the created operation.
-func ResourceSparseApplyAdagradDA(scope *Scope, var_ tf.Output, gradient_accumulator tf.Output, gradient_squared_accumulator tf.Output, grad tf.Output, indices tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, global_step tf.Output, optional ...ResourceSparseApplyAdagradDAAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceSparseApplyAdagradDA",
- Input: []tf.Input{
- var_, gradient_accumulator, gradient_squared_accumulator, grad, indices, lr, l1, l2, global_step,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// ResourceSparseApplyFtrlAttr is an optional argument to ResourceSparseApplyFtrl.
-type ResourceSparseApplyFtrlAttr func(optionalAttr)
-
-// ResourceSparseApplyFtrlUseLocking sets the optional use_locking attribute to value.
-//
-// value: If `True`, updating of the var and accum tensors will be protected
-// by a lock; otherwise the behavior is undefined, but may exhibit less
-// contention.
-// If not specified, defaults to false
-func ResourceSparseApplyFtrlUseLocking(value bool) ResourceSparseApplyFtrlAttr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// Update relevant entries in '*var' according to the Ftrl-proximal scheme.
-//
-// That is for rows we have grad for, we update var, accum and linear as follows:
-// accum_new = accum + grad * grad
-// linear += grad + (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
-// quadratic = 1.0 / (accum_new^(lr_power) * lr) + 2 * l2
-// var = (sign(linear) * l1 - linear) / quadratic if |linear| > l1 else 0.0
-// accum = accum_new
-//
-// Arguments:
-// var_: Should be from a Variable().
-// accum: Should be from a Variable().
-// linear: Should be from a Variable().
-// grad: The gradient.
-// indices: A vector of indices into the first dimension of var and accum.
-// lr: Scaling factor. Must be a scalar.
-// l1: L1 regularization. Must be a scalar.
-// l2: L2 regularization. Must be a scalar.
-// lr_power: Scaling factor. Must be a scalar.
-//
-// Returns the created operation.
-func ResourceSparseApplyFtrl(scope *Scope, var_ tf.Output, accum tf.Output, linear tf.Output, grad tf.Output, indices tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, lr_power tf.Output, optional ...ResourceSparseApplyFtrlAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceSparseApplyFtrl",
- Input: []tf.Input{
- var_, accum, linear, grad, indices, lr, l1, l2, lr_power,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// Returns which elements of x are Inf.
-//
-// @compatibility(numpy)
-// Equivalent to np.isinf
-// @end_compatibility
-func IsInf(scope *Scope, x tf.Output) (y tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "IsInf",
- Input: []tf.Input{
- x,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// TruncatedNormalAttr is an optional argument to TruncatedNormal.
-type TruncatedNormalAttr func(optionalAttr)
-
-// TruncatedNormalSeed sets the optional seed attribute to value.
-//
-// value: If either `seed` or `seed2` are set to be non-zero, the random number
-// generator is seeded by the given seed. Otherwise, it is seeded by a
-// random seed.
-// If not specified, defaults to 0
-func TruncatedNormalSeed(value int64) TruncatedNormalAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// TruncatedNormalSeed2 sets the optional seed2 attribute to value.
-//
-// value: A second seed to avoid seed collision.
-// If not specified, defaults to 0
-func TruncatedNormalSeed2(value int64) TruncatedNormalAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Outputs random values from a truncated normal distribution.
-//
-// The generated values follow a normal distribution with mean 0 and standard
-// deviation 1, except that values whose magnitude is more than 2 standard
-// deviations from the mean are dropped and re-picked.
-//
-// Arguments:
-// shape: The shape of the output tensor.
-// dtype: The type of the output.
-//
-// Returns A tensor of the specified shape filled with random truncated normal
-// values.
-func TruncatedNormal(scope *Scope, shape tf.Output, dtype tf.DataType, optional ...TruncatedNormalAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtype": dtype}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "TruncatedNormal",
- Input: []tf.Input{
- shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// SkipgramAttr is an optional argument to Skipgram.
-type SkipgramAttr func(optionalAttr)
-
-// SkipgramWindowSize sets the optional window_size attribute to value.
-//
-// value: The number of words to predict to the left and right of the target.
-// If not specified, defaults to 5
-func SkipgramWindowSize(value int64) SkipgramAttr {
- return func(m optionalAttr) {
- m["window_size"] = value
- }
-}
-
-// SkipgramMinCount sets the optional min_count attribute to value.
-//
-// value: The minimum number of word occurrences for it to be included in the
-// vocabulary.
-// If not specified, defaults to 5
-func SkipgramMinCount(value int64) SkipgramAttr {
- return func(m optionalAttr) {
- m["min_count"] = value
- }
-}
-
-// SkipgramSubsample sets the optional subsample attribute to value.
-//
-// value: Threshold for word occurrence. Words that appear with higher
-// frequency will be randomly down-sampled. Set to 0 to disable.
-// If not specified, defaults to 0.001
-func SkipgramSubsample(value float32) SkipgramAttr {
- return func(m optionalAttr) {
- m["subsample"] = value
- }
-}
-
-// Parses a text file and creates a batch of examples.
-//
-// DEPRECATED at GraphDef version 19: Moving word2vec into tensorflow_models/tutorials and deprecating its ops here as a result
-//
-// Arguments:
-// filename: The corpus's text file name.
-// batch_size: The size of produced batch.
-//
-// Returns A vector of words in the corpus.Frequencies of words. Sorted in the non-ascending order.Number of words per epoch in the data file.The current epoch number.The total number of words processed so far.A vector of word ids.A vector of word ids.
-func Skipgram(scope *Scope, filename string, batch_size int64, optional ...SkipgramAttr) (vocab_word tf.Output, vocab_freq tf.Output, words_per_epoch tf.Output, current_epoch tf.Output, total_words_processed tf.Output, examples tf.Output, labels tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"filename": filename, "batch_size": batch_size}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "Skipgram",
-
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2), op.Output(3), op.Output(4), op.Output(5), op.Output(6)
-}
-
-// StringToNumberAttr is an optional argument to StringToNumber.
-type StringToNumberAttr func(optionalAttr)
-
-// StringToNumberOutType sets the optional out_type attribute to value.
-//
-// value: The numeric type to interpret each string in `string_tensor` as.
-// If not specified, defaults to DT_FLOAT
-func StringToNumberOutType(value tf.DataType) StringToNumberAttr {
- return func(m optionalAttr) {
- m["out_type"] = value
- }
-}
-
-// Converts each string in the input Tensor to the specified numeric type.
-//
-// (Note that int32 overflow results in an error while float overflow
-// results in a rounded value.)
-//
-// Returns A Tensor of the same shape as the input `string_tensor`.
-func StringToNumber(scope *Scope, string_tensor tf.Output, optional ...StringToNumberAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "StringToNumber",
- Input: []tf.Input{
- string_tensor,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// ResourceApplyFtrlV2Attr is an optional argument to ResourceApplyFtrlV2.
-type ResourceApplyFtrlV2Attr func(optionalAttr)
-
-// ResourceApplyFtrlV2UseLocking sets the optional use_locking attribute to value.
-//
-// value: If `True`, updating of the var and accum tensors will be protected
-// by a lock; otherwise the behavior is undefined, but may exhibit less
-// contention.
-// If not specified, defaults to false
-func ResourceApplyFtrlV2UseLocking(value bool) ResourceApplyFtrlV2Attr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// Update '*var' according to the Ftrl-proximal scheme.
-//
-// grad_with_shrinkage = grad + 2 * l2_shrinkage * var
-// accum_new = accum + grad_with_shrinkage * grad_with_shrinkage
-// linear += grad_with_shrinkage +
-// (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
-// quadratic = 1.0 / (accum_new^(lr_power) * lr) + 2 * l2
-// var = (sign(linear) * l1 - linear) / quadratic if |linear| > l1 else 0.0
-// accum = accum_new
-//
-// Arguments:
-// var_: Should be from a Variable().
-// accum: Should be from a Variable().
-// linear: Should be from a Variable().
-// grad: The gradient.
-// lr: Scaling factor. Must be a scalar.
-// l1: L1 regulariation. Must be a scalar.
-// l2: L2 shrinkage regulariation. Must be a scalar.
-//
-// lr_power: Scaling factor. Must be a scalar.
-//
-// Returns the created operation.
-func ResourceApplyFtrlV2(scope *Scope, var_ tf.Output, accum tf.Output, linear tf.Output, grad tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, l2_shrinkage tf.Output, lr_power tf.Output, optional ...ResourceApplyFtrlV2Attr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceApplyFtrlV2",
- Input: []tf.Input{
- var_, accum, linear, grad, lr, l1, l2, l2_shrinkage, lr_power,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// Adds up a `SparseTensor` and a dense `Tensor`, producing a dense `Tensor`.
-//
-// This Op does not require `a_indices` be sorted in standard lexicographic order.
-//
-// Arguments:
-// a_indices: 2-D. The `indices` of the `SparseTensor`, with shape `[nnz, ndims]`.
-// a_values: 1-D. The `values` of the `SparseTensor`, with shape `[nnz]`.
-// a_shape: 1-D. The `shape` of the `SparseTensor`, with shape `[ndims]`.
-// b: `ndims`-D Tensor. With shape `a_shape`.
-func SparseTensorDenseAdd(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SparseTensorDenseAdd",
- Input: []tf.Input{
- a_indices, a_values, a_shape, b,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Encode audio data using the WAV file format.
-//
-// This operation will generate a string suitable to be saved out to create a .wav
-// audio file. It will be encoded in the 16-bit PCM format. It takes in float
-// values in the range -1.0f to 1.0f, and any outside that value will be clamped to
-// that range.
-//
-// `audio` is a 2-D float Tensor of shape `[length, channels]`.
-// `sample_rate` is a scalar Tensor holding the rate to use (e.g. 44100).
-//
-// Arguments:
-// audio: 2-D with shape `[length, channels]`.
-// sample_rate: Scalar containing the sample frequency.
-//
-// Returns 0-D. WAV-encoded file contents.
-func EncodeWav(scope *Scope, audio tf.Output, sample_rate tf.Output) (contents tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "EncodeWav",
- Input: []tf.Input{
- audio, sample_rate,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Computes atan of x element-wise.
-func Atan(scope *Scope, x tf.Output) (y tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Atan",
- Input: []tf.Input{
- x,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// ResourceApplyAdaMaxAttr is an optional argument to ResourceApplyAdaMax.
-type ResourceApplyAdaMaxAttr func(optionalAttr)
-
-// ResourceApplyAdaMaxUseLocking sets the optional use_locking attribute to value.
-//
-// value: If `True`, updating of the var, m, and v tensors will be protected
-// by a lock; otherwise the behavior is undefined, but may exhibit less
-// contention.
-// If not specified, defaults to false
-func ResourceApplyAdaMaxUseLocking(value bool) ResourceApplyAdaMaxAttr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// Update '*var' according to the AdaMax algorithm.
-//
-// m_t <- beta1 * m_{t-1} + (1 - beta1) * g
-// v_t <- max(beta2 * v_{t-1}, abs(g))
-// variable <- variable - learning_rate / (1 - beta1^t) * m_t / (v_t + epsilon)
-//
-// Arguments:
-// var_: Should be from a Variable().
-// m: Should be from a Variable().
-// v: Should be from a Variable().
-// beta1_power: Must be a scalar.
-// lr: Scaling factor. Must be a scalar.
-// beta1: Momentum factor. Must be a scalar.
-// beta2: Momentum factor. Must be a scalar.
-// epsilon: Ridge term. Must be a scalar.
-// grad: The gradient.
-//
-// Returns the created operation.
-func ResourceApplyAdaMax(scope *Scope, var_ tf.Output, m tf.Output, v tf.Output, beta1_power tf.Output, lr tf.Output, beta1 tf.Output, beta2 tf.Output, epsilon tf.Output, grad tf.Output, optional ...ResourceApplyAdaMaxAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceApplyAdaMax",
- Input: []tf.Input{
- var_, m, v, beta1_power, lr, beta1, beta2, epsilon, grad,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// AssertAttr is an optional argument to Assert.
-type AssertAttr func(optionalAttr)
-
-// AssertSummarize sets the optional summarize attribute to value.
-//
-// value: Print this many entries of each tensor.
-// If not specified, defaults to 3
-func AssertSummarize(value int64) AssertAttr {
- return func(m optionalAttr) {
- m["summarize"] = value
- }
-}
-
-// Asserts that the given condition is true.
-//
-// If `condition` evaluates to false, print the list of tensors in `data`.
-// `summarize` determines how many entries of the tensors to print.
-//
-// Arguments:
-// condition: The condition to evaluate.
-// data: The tensors to print out when condition is false.
-//
-// Returns the created operation.
-func Assert(scope *Scope, condition tf.Output, data []tf.Output, optional ...AssertAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "Assert",
- Input: []tf.Input{
- condition, tf.OutputList(data),
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// CudnnRNNBackpropAttr is an optional argument to CudnnRNNBackprop.
-type CudnnRNNBackpropAttr func(optionalAttr)
-
-// CudnnRNNBackpropRnnMode sets the optional rnn_mode attribute to value.
-// If not specified, defaults to "lstm"
-func CudnnRNNBackpropRnnMode(value string) CudnnRNNBackpropAttr {
- return func(m optionalAttr) {
- m["rnn_mode"] = value
- }
-}
-
-// CudnnRNNBackpropInputMode sets the optional input_mode attribute to value.
-// If not specified, defaults to "linear_input"
-func CudnnRNNBackpropInputMode(value string) CudnnRNNBackpropAttr {
- return func(m optionalAttr) {
- m["input_mode"] = value
- }
-}
-
-// CudnnRNNBackpropDirection sets the optional direction attribute to value.
-// If not specified, defaults to "unidirectional"
-func CudnnRNNBackpropDirection(value string) CudnnRNNBackpropAttr {
- return func(m optionalAttr) {
- m["direction"] = value
- }
-}
-
-// CudnnRNNBackpropDropout sets the optional dropout attribute to value.
-// If not specified, defaults to 0
-func CudnnRNNBackpropDropout(value float32) CudnnRNNBackpropAttr {
- return func(m optionalAttr) {
- m["dropout"] = value
- }
-}
-
-// CudnnRNNBackpropSeed sets the optional seed attribute to value.
-// If not specified, defaults to 0
-func CudnnRNNBackpropSeed(value int64) CudnnRNNBackpropAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// CudnnRNNBackpropSeed2 sets the optional seed2 attribute to value.
-// If not specified, defaults to 0
-func CudnnRNNBackpropSeed2(value int64) CudnnRNNBackpropAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Backprop step of CudnnRNN.
-//
-// Compute the backprop of both data and weights in a RNN.
-//
-// rnn_mode: Indicates the type of the RNN model.
-// input_mode: Indicate whether there is a linear projection between the input and
-// the actual computation before the first layer. 'skip_input' is only allowed
-// when input_size == num_units; 'auto_select' implies 'skip_input' when
-// input_size == num_units; otherwise, it implies 'linear_input'.
-// direction: Indicates whether a bidirectional model will be used. Should be
-// "unidirectional" or "bidirectional".
-// dropout: Dropout probability. When set to 0., dropout is disabled.
-// seed: The 1st part of a seed to initialize dropout.
-// seed2: The 2nd part of a seed to initialize dropout.
-// input: A 3-D tensor with the shape of [seq_length, batch_size, input_size].
-// input_h: A 3-D tensor with the shape of [num_layer * dir, batch_size,
-// num_units].
-// input_c: For LSTM, a 3-D tensor with the shape of
-// [num_layer * dir, batch, num_units]. For other models, it is ignored.
-// params: A 1-D tensor that contains the weights and biases in an opaque layout.
-// The size must be created through CudnnRNNParamsSize, and initialized
-// separately. Note that they might not be compatible across different
-// generations. So it is a good idea to save and restore
-// output: A 3-D tensor with the shape of [seq_length, batch_size,
-// dir * num_units].
-// output_h: The same shape has input_h.
-// output_c: The same shape as input_c for LSTM. An empty tensor for other models.
-// output_backprop: A 3-D tensor with the same shape as output in the forward pass.
-// output_h_backprop: A 3-D tensor with the same shape as output_h in the forward
-// pass.
-// output_c_backprop: A 3-D tensor with the same shape as output_c in the forward
-// pass.
-// reserve_space: The same reserve_space produced in for forward operation.
-// input_backprop: The backprop to input in the forward pass. Has the same shape
-// as input.
-// input_h_backprop: The backprop to input_h in the forward pass. Has the same
-// shape as input_h.
-// input_c_backprop: The backprop to input_c in the forward pass. Has the same
-// shape as input_c.
-// params_backprop: The backprop to the params buffer in the forward pass. Has the
-// same shape as params.
-func CudnnRNNBackprop(scope *Scope, input tf.Output, input_h tf.Output, input_c tf.Output, params tf.Output, output tf.Output, output_h tf.Output, output_c tf.Output, output_backprop tf.Output, output_h_backprop tf.Output, output_c_backprop tf.Output, reserve_space tf.Output, optional ...CudnnRNNBackpropAttr) (input_backprop tf.Output, input_h_backprop tf.Output, input_c_backprop tf.Output, params_backprop tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "CudnnRNNBackprop",
- Input: []tf.Input{
- input, input_h, input_c, params, output, output_h, output_c, output_backprop, output_h_backprop, output_c_backprop, reserve_space,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2), op.Output(3)
-}
-
-// Split a `SparseTensor` into `num_split` tensors along one dimension.
-//
-// If the `shape[split_dim]` is not an integer multiple of `num_split`. Slices
-// `[0 : shape[split_dim] % num_split]` gets one extra dimension.
-// For example, if `split_dim = 1` and `num_split = 2` and the input is
-//
-// input_tensor = shape = [2, 7]
-// [ a d e ]
-// [b c ]
-//
-// Graphically the output tensors are:
-//
-// output_tensor[0] = shape = [2, 4]
-// [ a ]
-// [b c ]
-//
-// output_tensor[1] = shape = [2, 3]
-// [ d e ]
-// [ ]
-//
-// Arguments:
-// split_dim: 0-D. The dimension along which to split. Must be in the range
-// `[0, rank(shape))`.
-// indices: 2-D tensor represents the indices of the sparse tensor.
-// values: 1-D tensor represents the values of the sparse tensor.
-// shape: 1-D. tensor represents the shape of the sparse tensor.
-// output indices: A list of 1-D tensors represents the indices of the output
-// sparse tensors.
-// num_split: The number of ways to split.
-//
-// Returns A list of 1-D tensors represents the values of the output sparse
-// tensors.A list of 1-D tensors represents the shape of the output sparse
-// tensors.
-func SparseSplit(scope *Scope, split_dim tf.Output, indices tf.Output, values tf.Output, shape tf.Output, num_split int64) (output_indices []tf.Output, output_values []tf.Output, output_shape []tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"num_split": num_split}
- opspec := tf.OpSpec{
- Type: "SparseSplit",
- Input: []tf.Input{
- split_dim, indices, values, shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- if scope.Err() != nil {
- return
- }
- var idx int
- var err error
- if output_indices, idx, err = makeOutputList(op, idx, "output_indices"); err != nil {
- scope.UpdateErr("SparseSplit", err)
- return
- }
- if output_values, idx, err = makeOutputList(op, idx, "output_values"); err != nil {
- scope.UpdateErr("SparseSplit", err)
- return
- }
- if output_shape, idx, err = makeOutputList(op, idx, "output_shape"); err != nil {
- scope.UpdateErr("SparseSplit", err)
- return
- }
- return output_indices, output_values, output_shape
-}
-
// Returns the element-wise sum of a list of tensors.
//
// `tf.accumulate_n_v2` performs the same operation as `tf.add_n`, but does not
@@ -13653,6 +13335,39 @@
return op.Output(0)
}
+// Adds up a SparseTensor and a dense Tensor, using these special rules:
+//
+// (1) Broadcasts the dense side to have the same shape as the sparse side, if
+// eligible;
+// (2) Then, only the dense values pointed to by the indices of the SparseTensor
+// participate in the cwise addition.
+//
+// By these rules, the result is a logical SparseTensor with exactly the same
+// indices and shape, but possibly with different non-zero values. The output of
+// this Op is the resultant non-zero values.
+//
+// Arguments:
+// sp_indices: 2-D. `N x R` matrix with the indices of non-empty values in a
+// SparseTensor, possibly not in canonical ordering.
+// sp_values: 1-D. `N` non-empty values corresponding to `sp_indices`.
+// sp_shape: 1-D. Shape of the input SparseTensor.
+// dense: `R`-D. The dense Tensor operand.
+//
+// Returns 1-D. The `N` values that are operated on.
+func SparseDenseCwiseAdd(scope *Scope, sp_indices tf.Output, sp_values tf.Output, sp_shape tf.Output, dense tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseDenseCwiseAdd",
+ Input: []tf.Input{
+ sp_indices, sp_values, sp_shape, dense,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Computes the complementary error function of `x` element-wise.
func Erfc(scope *Scope, x tf.Output) (y tf.Output) {
if scope.Err() != nil {
@@ -14066,52 +13781,6 @@
return op.Output(0)
}
-// MaxPoolAttr is an optional argument to MaxPool.
-type MaxPoolAttr func(optionalAttr)
-
-// MaxPoolDataFormat sets the optional data_format attribute to value.
-//
-// value: Specify the data format of the input and output data. With the
-// default format "NHWC", the data is stored in the order of:
-// [batch, in_height, in_width, in_channels].
-// Alternatively, the format could be "NCHW", the data storage order of:
-// [batch, in_channels, in_height, in_width].
-// If not specified, defaults to "NHWC"
-func MaxPoolDataFormat(value string) MaxPoolAttr {
- return func(m optionalAttr) {
- m["data_format"] = value
- }
-}
-
-// Performs max pooling on the input.
-//
-// Arguments:
-// input: 4-D input to pool over.
-// ksize: The size of the window for each dimension of the input tensor.
-// strides: The stride of the sliding window for each dimension of the
-// input tensor.
-// padding: The type of padding algorithm to use.
-//
-// Returns The max pooled output tensor.
-func MaxPool(scope *Scope, input tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPoolAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MaxPool",
- Input: []tf.Input{
- input,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Concatenates quantized tensors along one dimension.
//
// Arguments:
@@ -14249,66 +13918,6 @@
return op.Output(0), op.Output(1), op.Output(2)
}
-// Returns the truth value of (x != y) element-wise.
-//
-// *NOTE*: `NotEqual` supports broadcasting. More about broadcasting
-// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
-func NotEqual(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "NotEqual",
- Input: []tf.Input{
- x, y,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Inverse 3D real-valued fast Fourier transform.
-//
-// Computes the inverse 3-dimensional discrete Fourier transform of a real-valued
-// signal over the inner-most 3 dimensions of `input`.
-//
-// The inner-most 3 dimensions of `input` are assumed to be the result of `RFFT3D`:
-// The inner-most dimension contains the `fft_length / 2 + 1` unique components of
-// the DFT of a real-valued signal. If `fft_length` is not provided, it is computed
-// from the size of the inner-most 3 dimensions of `input`. If the FFT length used
-// to compute `input` is odd, it should be provided since it cannot be inferred
-// properly.
-//
-// Along each axis `IRFFT3D` is computed on, if `fft_length` (or
-// `fft_length / 2 + 1` for the inner-most dimension) is smaller than the
-// corresponding dimension of `input`, the dimension is cropped. If it is larger,
-// the dimension is padded with zeros.
-//
-// Arguments:
-// input: A complex64 tensor.
-// fft_length: An int32 tensor of shape [3]. The FFT length for each dimension.
-//
-// Returns A float32 tensor of the same rank as `input`. The inner-most 3
-// dimensions of `input` are replaced with the `fft_length` samples of their
-// inverse 3D real Fourier transform.
-//
-// @compatibility(numpy)
-// Equivalent to np.irfftn with 3 dimensions.
-// @end_compatibility
-func IRFFT3D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "IRFFT3D",
- Input: []tf.Input{
- input, fft_length,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// StringSplitAttr is an optional argument to StringSplit.
type StringSplitAttr func(optionalAttr)
@@ -14463,176 +14072,99 @@
return op.Output(0)
}
-// AvgPoolGradAttr is an optional argument to AvgPoolGrad.
-type AvgPoolGradAttr func(optionalAttr)
+// CudnnRNNBackpropAttr is an optional argument to CudnnRNNBackprop.
+type CudnnRNNBackpropAttr func(optionalAttr)
-// AvgPoolGradDataFormat sets the optional data_format attribute to value.
-//
-// value: Specify the data format of the input and output data. With the
-// default format "NHWC", the data is stored in the order of:
-// [batch, in_height, in_width, in_channels].
-// Alternatively, the format could be "NCHW", the data storage order of:
-// [batch, in_channels, in_height, in_width].
-// If not specified, defaults to "NHWC"
-func AvgPoolGradDataFormat(value string) AvgPoolGradAttr {
+// CudnnRNNBackpropRnnMode sets the optional rnn_mode attribute to value.
+// If not specified, defaults to "lstm"
+func CudnnRNNBackpropRnnMode(value string) CudnnRNNBackpropAttr {
return func(m optionalAttr) {
- m["data_format"] = value
+ m["rnn_mode"] = value
}
}
-// Computes gradients of the average pooling function.
-//
-// Arguments:
-// orig_input_shape: 1-D. Shape of the original input to `avg_pool`.
-// grad: 4-D with shape `[batch, height, width, channels]`. Gradients w.r.t.
-// the output of `avg_pool`.
-// ksize: The size of the sliding window for each dimension of the input.
-// strides: The stride of the sliding window for each dimension of the input.
-// padding: The type of padding algorithm to use.
-//
-// Returns 4-D. Gradients w.r.t. the input of `avg_pool`.
-func AvgPoolGrad(scope *Scope, orig_input_shape tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...AvgPoolGradAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
+// CudnnRNNBackpropInputMode sets the optional input_mode attribute to value.
+// If not specified, defaults to "linear_input"
+func CudnnRNNBackpropInputMode(value string) CudnnRNNBackpropAttr {
+ return func(m optionalAttr) {
+ m["input_mode"] = value
}
- attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "AvgPoolGrad",
- Input: []tf.Input{
- orig_input_shape, grad,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
}
-// Greedily selects a subset of bounding boxes in descending order of score,
-//
-// pruning away boxes that have high overlaps
-// with previously selected boxes. Bounding boxes with score less than
-// `score_threshold` are removed. N-by-n overlap values are supplied as square matrix,
-// which allows for defining a custom overlap criterium (eg. intersection over union,
-// intersection over area, etc.).
-//
-// The output of this operation is a set of integers indexing into the input
-// collection of bounding boxes representing the selected boxes. The bounding
-// box coordinates corresponding to the selected indices can then be obtained
-// using the `tf.gather operation`. For example:
-//
-// selected_indices = tf.image.non_max_suppression_with_overlaps(
-// overlaps, scores, max_output_size, overlap_threshold, score_threshold)
-// selected_boxes = tf.gather(boxes, selected_indices)
-//
-// Arguments:
-// overlaps: A 2-D float tensor of shape `[num_boxes, num_boxes]` representing
-// the n-by-n box overlap values.
-// scores: A 1-D float tensor of shape `[num_boxes]` representing a single
-// score corresponding to each box (each row of boxes).
-// max_output_size: A scalar integer tensor representing the maximum number of
-// boxes to be selected by non max suppression.
-// overlap_threshold: A 0-D float tensor representing the threshold for deciding whether
-// boxes overlap too.
-// score_threshold: A 0-D float tensor representing the threshold for deciding when to remove
-// boxes based on score.
-//
-// Returns A 1-D integer tensor of shape `[M]` representing the selected
-// indices from the boxes tensor, where `M <= max_output_size`.
-func NonMaxSuppressionWithOverlaps(scope *Scope, overlaps tf.Output, scores tf.Output, max_output_size tf.Output, overlap_threshold tf.Output, score_threshold tf.Output) (selected_indices tf.Output) {
- if scope.Err() != nil {
- return
+// CudnnRNNBackpropDirection sets the optional direction attribute to value.
+// If not specified, defaults to "unidirectional"
+func CudnnRNNBackpropDirection(value string) CudnnRNNBackpropAttr {
+ return func(m optionalAttr) {
+ m["direction"] = value
}
- opspec := tf.OpSpec{
- Type: "NonMaxSuppressionWithOverlaps",
- Input: []tf.Input{
- overlaps, scores, max_output_size, overlap_threshold, score_threshold,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
}
-// StageClearAttr is an optional argument to StageClear.
-type StageClearAttr func(optionalAttr)
-
-// StageClearCapacity sets the optional capacity attribute to value.
+// CudnnRNNBackpropDropout sets the optional dropout attribute to value.
// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func StageClearCapacity(value int64) StageClearAttr {
+func CudnnRNNBackpropDropout(value float32) CudnnRNNBackpropAttr {
return func(m optionalAttr) {
- m["capacity"] = value
+ m["dropout"] = value
}
}
-// StageClearMemoryLimit sets the optional memory_limit attribute to value.
+// CudnnRNNBackpropSeed sets the optional seed attribute to value.
// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func StageClearMemoryLimit(value int64) StageClearAttr {
+func CudnnRNNBackpropSeed(value int64) CudnnRNNBackpropAttr {
return func(m optionalAttr) {
- m["memory_limit"] = value
+ m["seed"] = value
}
}
-// StageClearContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func StageClearContainer(value string) StageClearAttr {
+// CudnnRNNBackpropSeed2 sets the optional seed2 attribute to value.
+// If not specified, defaults to 0
+func CudnnRNNBackpropSeed2(value int64) CudnnRNNBackpropAttr {
return func(m optionalAttr) {
- m["container"] = value
+ m["seed2"] = value
}
}
-// StageClearSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func StageClearSharedName(value string) StageClearAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op removes all elements in the underlying container.
+// Backprop step of CudnnRNN.
//
-// Returns the created operation.
-func StageClear(scope *Scope, dtypes []tf.DataType, optional ...StageClearAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "StageClear",
-
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// QuantizedRelu6Attr is an optional argument to QuantizedRelu6.
-type QuantizedRelu6Attr func(optionalAttr)
-
-// QuantizedRelu6OutType sets the optional out_type attribute to value.
-// If not specified, defaults to DT_QUINT8
-func QuantizedRelu6OutType(value tf.DataType) QuantizedRelu6Attr {
- return func(m optionalAttr) {
- m["out_type"] = value
- }
-}
-
-// Computes Quantized Rectified Linear 6: `min(max(features, 0), 6)`
+// Compute the backprop of both data and weights in a RNN.
//
-// Arguments:
-//
-// min_features: The float value that the lowest quantized value represents.
-// max_features: The float value that the highest quantized value represents.
-//
-// Returns Has the same output shape as "features".The float value that the lowest quantized value represents.The float value that the highest quantized value represents.
-func QuantizedRelu6(scope *Scope, features tf.Output, min_features tf.Output, max_features tf.Output, optional ...QuantizedRelu6Attr) (activations tf.Output, min_activations tf.Output, max_activations tf.Output) {
+// rnn_mode: Indicates the type of the RNN model.
+// input_mode: Indicate whether there is a linear projection between the input and
+// the actual computation before the first layer. 'skip_input' is only allowed
+// when input_size == num_units; 'auto_select' implies 'skip_input' when
+// input_size == num_units; otherwise, it implies 'linear_input'.
+// direction: Indicates whether a bidirectional model will be used. Should be
+// "unidirectional" or "bidirectional".
+// dropout: Dropout probability. When set to 0., dropout is disabled.
+// seed: The 1st part of a seed to initialize dropout.
+// seed2: The 2nd part of a seed to initialize dropout.
+// input: A 3-D tensor with the shape of [seq_length, batch_size, input_size].
+// input_h: A 3-D tensor with the shape of [num_layer * dir, batch_size,
+// num_units].
+// input_c: For LSTM, a 3-D tensor with the shape of
+// [num_layer * dir, batch, num_units]. For other models, it is ignored.
+// params: A 1-D tensor that contains the weights and biases in an opaque layout.
+// The size must be created through CudnnRNNParamsSize, and initialized
+// separately. Note that they might not be compatible across different
+// generations. So it is a good idea to save and restore
+// output: A 3-D tensor with the shape of [seq_length, batch_size,
+// dir * num_units].
+// output_h: The same shape has input_h.
+// output_c: The same shape as input_c for LSTM. An empty tensor for other models.
+// output_backprop: A 3-D tensor with the same shape as output in the forward pass.
+// output_h_backprop: A 3-D tensor with the same shape as output_h in the forward
+// pass.
+// output_c_backprop: A 3-D tensor with the same shape as output_c in the forward
+// pass.
+// reserve_space: The same reserve_space produced in for forward operation.
+// input_backprop: The backprop to input in the forward pass. Has the same shape
+// as input.
+// input_h_backprop: The backprop to input_h in the forward pass. Has the same
+// shape as input_h.
+// input_c_backprop: The backprop to input_c in the forward pass. Has the same
+// shape as input_c.
+// params_backprop: The backprop to the params buffer in the forward pass. Has the
+// same shape as params.
+func CudnnRNNBackprop(scope *Scope, input tf.Output, input_h tf.Output, input_c tf.Output, params tf.Output, output tf.Output, output_h tf.Output, output_c tf.Output, output_backprop tf.Output, output_h_backprop tf.Output, output_c_backprop tf.Output, reserve_space tf.Output, optional ...CudnnRNNBackpropAttr) (input_backprop tf.Output, input_h_backprop tf.Output, input_c_backprop tf.Output, params_backprop tf.Output) {
if scope.Err() != nil {
return
}
@@ -14641,104 +14173,216 @@
a(attrs)
}
opspec := tf.OpSpec{
- Type: "QuantizedRelu6",
+ Type: "CudnnRNNBackprop",
Input: []tf.Input{
- features, min_features, max_features,
+ input, input_h, input_c, params, output, output_h, output_c, output_backprop, output_h_backprop, output_c_backprop, reserve_space,
},
Attrs: attrs,
}
op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
+ return op.Output(0), op.Output(1), op.Output(2), op.Output(3)
}
-// FixedLengthRecordReaderV2Attr is an optional argument to FixedLengthRecordReaderV2.
-type FixedLengthRecordReaderV2Attr func(optionalAttr)
-
-// FixedLengthRecordReaderV2HeaderBytes sets the optional header_bytes attribute to value.
+// Encode audio data using the WAV file format.
//
-// value: Number of bytes in the header, defaults to 0.
-// If not specified, defaults to 0
-func FixedLengthRecordReaderV2HeaderBytes(value int64) FixedLengthRecordReaderV2Attr {
- return func(m optionalAttr) {
- m["header_bytes"] = value
- }
-}
-
-// FixedLengthRecordReaderV2FooterBytes sets the optional footer_bytes attribute to value.
+// This operation will generate a string suitable to be saved out to create a .wav
+// audio file. It will be encoded in the 16-bit PCM format. It takes in float
+// values in the range -1.0f to 1.0f, and any outside that value will be clamped to
+// that range.
//
-// value: Number of bytes in the footer, defaults to 0.
-// If not specified, defaults to 0
-func FixedLengthRecordReaderV2FooterBytes(value int64) FixedLengthRecordReaderV2Attr {
- return func(m optionalAttr) {
- m["footer_bytes"] = value
- }
-}
-
-// FixedLengthRecordReaderV2HopBytes sets the optional hop_bytes attribute to value.
-//
-// value: Number of bytes to hop before each read. Default of 0 means using
-// record_bytes.
-// If not specified, defaults to 0
-func FixedLengthRecordReaderV2HopBytes(value int64) FixedLengthRecordReaderV2Attr {
- return func(m optionalAttr) {
- m["hop_bytes"] = value
- }
-}
-
-// FixedLengthRecordReaderV2Container sets the optional container attribute to value.
-//
-// value: If non-empty, this reader is placed in the given container.
-// Otherwise, a default container is used.
-// If not specified, defaults to ""
-func FixedLengthRecordReaderV2Container(value string) FixedLengthRecordReaderV2Attr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// FixedLengthRecordReaderV2SharedName sets the optional shared_name attribute to value.
-//
-// value: If non-empty, this reader is named in the given bucket
-// with this shared_name. Otherwise, the node name is used instead.
-// If not specified, defaults to ""
-func FixedLengthRecordReaderV2SharedName(value string) FixedLengthRecordReaderV2Attr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// FixedLengthRecordReaderV2Encoding sets the optional encoding attribute to value.
-//
-// value: The type of encoding for the file. Currently ZLIB and GZIP
-// are supported. Defaults to none.
-// If not specified, defaults to ""
-func FixedLengthRecordReaderV2Encoding(value string) FixedLengthRecordReaderV2Attr {
- return func(m optionalAttr) {
- m["encoding"] = value
- }
-}
-
-// A Reader that outputs fixed-length records from a file.
+// `audio` is a 2-D float Tensor of shape `[length, channels]`.
+// `sample_rate` is a scalar Tensor holding the rate to use (e.g. 44100).
//
// Arguments:
-// record_bytes: Number of bytes in the record.
+// audio: 2-D with shape `[length, channels]`.
+// sample_rate: Scalar containing the sample frequency.
//
-// Returns The handle to reference the Reader.
-func FixedLengthRecordReaderV2(scope *Scope, record_bytes int64, optional ...FixedLengthRecordReaderV2Attr) (reader_handle tf.Output) {
+// Returns 0-D. WAV-encoded file contents.
+func EncodeWav(scope *Scope, audio tf.Output, sample_rate tf.Output) (contents tf.Output) {
if scope.Err() != nil {
return
}
- attrs := map[string]interface{}{"record_bytes": record_bytes}
+ opspec := tf.OpSpec{
+ Type: "EncodeWav",
+ Input: []tf.Input{
+ audio, sample_rate,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Computes atan of x element-wise.
+func Atan(scope *Scope, x tf.Output) (y tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Atan",
+ Input: []tf.Input{
+ x,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ResourceApplyAdaMaxAttr is an optional argument to ResourceApplyAdaMax.
+type ResourceApplyAdaMaxAttr func(optionalAttr)
+
+// ResourceApplyAdaMaxUseLocking sets the optional use_locking attribute to value.
+//
+// value: If `True`, updating of the var, m, and v tensors will be protected
+// by a lock; otherwise the behavior is undefined, but may exhibit less
+// contention.
+// If not specified, defaults to false
+func ResourceApplyAdaMaxUseLocking(value bool) ResourceApplyAdaMaxAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Update '*var' according to the AdaMax algorithm.
+//
+// m_t <- beta1 * m_{t-1} + (1 - beta1) * g
+// v_t <- max(beta2 * v_{t-1}, abs(g))
+// variable <- variable - learning_rate / (1 - beta1^t) * m_t / (v_t + epsilon)
+//
+// Arguments:
+// var_: Should be from a Variable().
+// m: Should be from a Variable().
+// v: Should be from a Variable().
+// beta1_power: Must be a scalar.
+// lr: Scaling factor. Must be a scalar.
+// beta1: Momentum factor. Must be a scalar.
+// beta2: Momentum factor. Must be a scalar.
+// epsilon: Ridge term. Must be a scalar.
+// grad: The gradient.
+//
+// Returns the created operation.
+func ResourceApplyAdaMax(scope *Scope, var_ tf.Output, m tf.Output, v tf.Output, beta1_power tf.Output, lr tf.Output, beta1 tf.Output, beta2 tf.Output, epsilon tf.Output, grad tf.Output, optional ...ResourceApplyAdaMaxAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
for _, a := range optional {
a(attrs)
}
opspec := tf.OpSpec{
- Type: "FixedLengthRecordReaderV2",
+ Type: "ResourceApplyAdaMax",
+ Input: []tf.Input{
+ var_, m, v, beta1_power, lr, beta1, beta2, epsilon, grad,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+// AssertAttr is an optional argument to Assert.
+type AssertAttr func(optionalAttr)
+
+// AssertSummarize sets the optional summarize attribute to value.
+//
+// value: Print this many entries of each tensor.
+// If not specified, defaults to 3
+func AssertSummarize(value int64) AssertAttr {
+ return func(m optionalAttr) {
+ m["summarize"] = value
+ }
+}
+
+// Asserts that the given condition is true.
+//
+// If `condition` evaluates to false, print the list of tensors in `data`.
+// `summarize` determines how many entries of the tensors to print.
+//
+// Arguments:
+// condition: The condition to evaluate.
+// data: The tensors to print out when condition is false.
+//
+// Returns the created operation.
+func Assert(scope *Scope, condition tf.Output, data []tf.Output, optional ...AssertAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "Assert",
+ Input: []tf.Input{
+ condition, tf.OutputList(data),
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Split a `SparseTensor` into `num_split` tensors along one dimension.
+//
+// If the `shape[split_dim]` is not an integer multiple of `num_split`. Slices
+// `[0 : shape[split_dim] % num_split]` gets one extra dimension.
+// For example, if `split_dim = 1` and `num_split = 2` and the input is
+//
+// input_tensor = shape = [2, 7]
+// [ a d e ]
+// [b c ]
+//
+// Graphically the output tensors are:
+//
+// output_tensor[0] = shape = [2, 4]
+// [ a ]
+// [b c ]
+//
+// output_tensor[1] = shape = [2, 3]
+// [ d e ]
+// [ ]
+//
+// Arguments:
+// split_dim: 0-D. The dimension along which to split. Must be in the range
+// `[0, rank(shape))`.
+// indices: 2-D tensor represents the indices of the sparse tensor.
+// values: 1-D tensor represents the values of the sparse tensor.
+// shape: 1-D. tensor represents the shape of the sparse tensor.
+// output indices: A list of 1-D tensors represents the indices of the output
+// sparse tensors.
+// num_split: The number of ways to split.
+//
+// Returns A list of 1-D tensors represents the values of the output sparse
+// tensors.A list of 1-D tensors represents the shape of the output sparse
+// tensors.
+func SparseSplit(scope *Scope, split_dim tf.Output, indices tf.Output, values tf.Output, shape tf.Output, num_split int64) (output_indices []tf.Output, output_values []tf.Output, output_shape []tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"num_split": num_split}
+ opspec := tf.OpSpec{
+ Type: "SparseSplit",
+ Input: []tf.Input{
+ split_dim, indices, values, shape,
+ },
Attrs: attrs,
}
op := scope.AddOperation(opspec)
- return op.Output(0)
+ if scope.Err() != nil {
+ return
+ }
+ var idx int
+ var err error
+ if output_indices, idx, err = makeOutputList(op, idx, "output_indices"); err != nil {
+ scope.UpdateErr("SparseSplit", err)
+ return
+ }
+ if output_values, idx, err = makeOutputList(op, idx, "output_values"); err != nil {
+ scope.UpdateErr("SparseSplit", err)
+ return
+ }
+ if output_shape, idx, err = makeOutputList(op, idx, "output_shape"); err != nil {
+ scope.UpdateErr("SparseSplit", err)
+ return
+ }
+ return output_indices, output_values, output_shape
}
// Computes numerical negative value element-wise.
@@ -15437,355 +15081,6 @@
return op.Output(0)
}
-// The gradient operator for the SparseAdd op.
-//
-// The SparseAdd op calculates A + B, where A, B, and the sum are all represented
-// as `SparseTensor` objects. This op takes in the upstream gradient w.r.t.
-// non-empty values of the sum, and outputs the gradients w.r.t. the non-empty
-// values of A and B.
-//
-// Arguments:
-// backprop_val_grad: 1-D with shape `[nnz(sum)]`. The gradient with respect to
-// the non-empty values of the sum.
-// a_indices: 2-D. The `indices` of the `SparseTensor` A, size `[nnz(A), ndims]`.
-// b_indices: 2-D. The `indices` of the `SparseTensor` B, size `[nnz(B), ndims]`.
-// sum_indices: 2-D. The `indices` of the sum `SparseTensor`, size
-// `[nnz(sum), ndims]`.
-//
-// Returns 1-D with shape `[nnz(A)]`. The gradient with respect to the
-// non-empty values of A.1-D with shape `[nnz(B)]`. The gradient with respect to the
-// non-empty values of B.
-func SparseAddGrad(scope *Scope, backprop_val_grad tf.Output, a_indices tf.Output, b_indices tf.Output, sum_indices tf.Output) (a_val_grad tf.Output, b_val_grad tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SparseAddGrad",
- Input: []tf.Input{
- backprop_val_grad, a_indices, b_indices, sum_indices,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1)
-}
-
-// ResourceSparseApplyCenteredRMSPropAttr is an optional argument to ResourceSparseApplyCenteredRMSProp.
-type ResourceSparseApplyCenteredRMSPropAttr func(optionalAttr)
-
-// ResourceSparseApplyCenteredRMSPropUseLocking sets the optional use_locking attribute to value.
-//
-// value: If `True`, updating of the var, mg, ms, and mom tensors is
-// protected by a lock; otherwise the behavior is undefined, but may exhibit less
-// contention.
-// If not specified, defaults to false
-func ResourceSparseApplyCenteredRMSPropUseLocking(value bool) ResourceSparseApplyCenteredRMSPropAttr {
- return func(m optionalAttr) {
- m["use_locking"] = value
- }
-}
-
-// Update '*var' according to the centered RMSProp algorithm.
-//
-// 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)
-//
-// ms <- rho * ms_{t-1} + (1-rho) * grad * grad
-// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms + epsilon)
-// var <- var - mom
-//
-// Arguments:
-// var_: Should be from a Variable().
-// mg: Should be from a Variable().
-// ms: Should be from a Variable().
-// mom: Should be from a Variable().
-// lr: Scaling factor. Must be a scalar.
-// rho: Decay rate. Must be a scalar.
-//
-// epsilon: Ridge term. Must be a scalar.
-// grad: The gradient.
-// indices: A vector of indices into the first dimension of var, ms and mom.
-//
-// Returns the created operation.
-func ResourceSparseApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms tf.Output, mom tf.Output, lr tf.Output, rho tf.Output, momentum tf.Output, epsilon tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyCenteredRMSPropAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ResourceSparseApplyCenteredRMSProp",
- Input: []tf.Input{
- var_, mg, ms, mom, lr, rho, momentum, epsilon, grad, indices,
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// Creates a dataset that batches `batch_size` elements from `input_dataset`.
-//
-// Arguments:
-//
-// batch_size: A scalar representing the number of elements to accumulate in a
-// batch.
-//
-//
-func BatchDataset(scope *Scope, input_dataset tf.Output, batch_size tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes}
- opspec := tf.OpSpec{
- Type: "BatchDataset",
- Input: []tf.Input{
- input_dataset, batch_size,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// RandomPoissonV2Attr is an optional argument to RandomPoissonV2.
-type RandomPoissonV2Attr func(optionalAttr)
-
-// RandomPoissonV2Seed sets the optional seed attribute to value.
-//
-// value: If either `seed` or `seed2` are set to be non-zero, the random number
-// generator is seeded by the given seed. Otherwise, it is seeded by a
-// random seed.
-// If not specified, defaults to 0
-func RandomPoissonV2Seed(value int64) RandomPoissonV2Attr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// RandomPoissonV2Seed2 sets the optional seed2 attribute to value.
-//
-// value: A second seed to avoid seed collision.
-// If not specified, defaults to 0
-func RandomPoissonV2Seed2(value int64) RandomPoissonV2Attr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// RandomPoissonV2Dtype sets the optional dtype attribute to value.
-// If not specified, defaults to DT_INT64
-func RandomPoissonV2Dtype(value tf.DataType) RandomPoissonV2Attr {
- return func(m optionalAttr) {
- m["dtype"] = value
- }
-}
-
-// Outputs random values from the Poisson distribution(s) described by rate.
-//
-// This op uses two algorithms, depending on rate. If rate >= 10, then
-// the algorithm by Hormann is used to acquire samples via
-// transformation-rejection.
-// See http://www.sciencedirect.com/science/article/pii/0167668793909974.
-//
-// Otherwise, Knuth's algorithm is used to acquire samples via multiplying uniform
-// random variables.
-// See Donald E. Knuth (1969). Seminumerical Algorithms. The Art of Computer
-// Programming, Volume 2. Addison Wesley
-//
-// Arguments:
-// shape: 1-D integer tensor. Shape of independent samples to draw from each
-// distribution described by the shape parameters given in rate.
-// rate: A tensor in which each scalar is a "rate" parameter describing the
-// associated poisson distribution.
-//
-// Returns A tensor with shape `shape + shape(rate)`. Each slice
-// `[:, ..., :, i0, i1, ...iN]` contains the samples drawn for
-// `rate[i0, i1, ...iN]`.
-func RandomPoissonV2(scope *Scope, shape tf.Output, rate tf.Output, optional ...RandomPoissonV2Attr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "RandomPoissonV2",
- Input: []tf.Input{
- shape, rate,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// DecodeAndCropJpegAttr is an optional argument to DecodeAndCropJpeg.
-type DecodeAndCropJpegAttr func(optionalAttr)
-
-// DecodeAndCropJpegChannels sets the optional channels attribute to value.
-//
-// value: Number of color channels for the decoded image.
-// If not specified, defaults to 0
-func DecodeAndCropJpegChannels(value int64) DecodeAndCropJpegAttr {
- return func(m optionalAttr) {
- m["channels"] = value
- }
-}
-
-// DecodeAndCropJpegRatio sets the optional ratio attribute to value.
-//
-// value: Downscaling ratio.
-// If not specified, defaults to 1
-func DecodeAndCropJpegRatio(value int64) DecodeAndCropJpegAttr {
- return func(m optionalAttr) {
- m["ratio"] = value
- }
-}
-
-// DecodeAndCropJpegFancyUpscaling sets the optional fancy_upscaling attribute to value.
-//
-// value: If true use a slower but nicer upscaling of the
-// chroma planes (yuv420/422 only).
-// If not specified, defaults to true
-func DecodeAndCropJpegFancyUpscaling(value bool) DecodeAndCropJpegAttr {
- return func(m optionalAttr) {
- m["fancy_upscaling"] = value
- }
-}
-
-// DecodeAndCropJpegTryRecoverTruncated sets the optional try_recover_truncated attribute to value.
-//
-// value: If true try to recover an image from truncated input.
-// If not specified, defaults to false
-func DecodeAndCropJpegTryRecoverTruncated(value bool) DecodeAndCropJpegAttr {
- return func(m optionalAttr) {
- m["try_recover_truncated"] = value
- }
-}
-
-// DecodeAndCropJpegAcceptableFraction sets the optional acceptable_fraction attribute to value.
-//
-// value: The minimum required fraction of lines before a truncated
-// input is accepted.
-// If not specified, defaults to 1
-func DecodeAndCropJpegAcceptableFraction(value float32) DecodeAndCropJpegAttr {
- return func(m optionalAttr) {
- m["acceptable_fraction"] = value
- }
-}
-
-// DecodeAndCropJpegDctMethod sets the optional dct_method attribute to value.
-//
-// value: string specifying a hint about the algorithm used for
-// decompression. Defaults to "" which maps to a system-specific
-// default. Currently valid values are ["INTEGER_FAST",
-// "INTEGER_ACCURATE"]. The hint may be ignored (e.g., the internal
-// jpeg library changes to a version that does not have that specific
-// option.)
-// If not specified, defaults to ""
-func DecodeAndCropJpegDctMethod(value string) DecodeAndCropJpegAttr {
- return func(m optionalAttr) {
- m["dct_method"] = value
- }
-}
-
-// Decode and Crop a JPEG-encoded image to a uint8 tensor.
-//
-// The attr `channels` indicates the desired number of color channels for the
-// decoded image.
-//
-// Accepted values are:
-//
-// * 0: Use the number of channels in the JPEG-encoded image.
-// * 1: output a grayscale image.
-// * 3: output an RGB image.
-//
-// If needed, the JPEG-encoded image is transformed to match the requested number
-// of color channels.
-//
-// The attr `ratio` allows downscaling the image by an integer factor during
-// decoding. Allowed values are: 1, 2, 4, and 8. This is much faster than
-// downscaling the image later.
-//
-//
-// It is equivalent to a combination of decode and crop, but much faster by only
-// decoding partial jpeg image.
-//
-// Arguments:
-// contents: 0-D. The JPEG-encoded image.
-// crop_window: 1-D. The crop window: [crop_y, crop_x, crop_height, crop_width].
-//
-// Returns 3-D with shape `[height, width, channels]`..
-func DecodeAndCropJpeg(scope *Scope, contents tf.Output, crop_window tf.Output, optional ...DecodeAndCropJpegAttr) (image tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "DecodeAndCropJpeg",
- Input: []tf.Input{
- contents, crop_window,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Adds two `SparseTensor` objects to produce another `SparseTensor`.
-//
-// The input `SparseTensor` objects' indices are assumed ordered in standard
-// lexicographic order. If this is not the case, before this step run
-// `SparseReorder` to restore index ordering.
-//
-// By default, if two values sum to zero at some index, the output `SparseTensor`
-// would still include that particular location in its index, storing a zero in the
-// corresponding value slot. To override this, callers can specify `thresh`,
-// indicating that if the sum has a magnitude strictly smaller than `thresh`, its
-// corresponding value and index would then not be included. In particular,
-// `thresh == 0` (default) means everything is kept and actual thresholding happens
-// only for a positive value.
-//
-// In the following shapes, `nnz` is the count after taking `thresh` into account.
-//
-// Arguments:
-// a_indices: 2-D. The `indices` of the first `SparseTensor`, size `[nnz, ndims]` Matrix.
-// a_values: 1-D. The `values` of the first `SparseTensor`, size `[nnz]` Vector.
-// a_shape: 1-D. The `shape` of the first `SparseTensor`, size `[ndims]` Vector.
-// b_indices: 2-D. The `indices` of the second `SparseTensor`, size `[nnz, ndims]` Matrix.
-// b_values: 1-D. The `values` of the second `SparseTensor`, size `[nnz]` Vector.
-// b_shape: 1-D. The `shape` of the second `SparseTensor`, size `[ndims]` Vector.
-// thresh: 0-D. The magnitude threshold that determines if an output value/index
-// pair takes space.
-func SparseAdd(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b_indices tf.Output, b_values tf.Output, b_shape tf.Output, thresh tf.Output) (sum_indices tf.Output, sum_values tf.Output, sum_shape tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SparseAdd",
- Input: []tf.Input{
- a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// LRNAttr is an optional argument to LRN.
type LRNAttr func(optionalAttr)
@@ -16455,6 +15750,531 @@
return op.Output(0)
}
+// ResourceSparseApplyFtrlAttr is an optional argument to ResourceSparseApplyFtrl.
+type ResourceSparseApplyFtrlAttr func(optionalAttr)
+
+// ResourceSparseApplyFtrlUseLocking sets the optional use_locking attribute to value.
+//
+// value: If `True`, updating of the var and accum tensors will be protected
+// by a lock; otherwise the behavior is undefined, but may exhibit less
+// contention.
+// If not specified, defaults to false
+func ResourceSparseApplyFtrlUseLocking(value bool) ResourceSparseApplyFtrlAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Update relevant entries in '*var' according to the Ftrl-proximal scheme.
+//
+// That is for rows we have grad for, we update var, accum and linear as follows:
+// accum_new = accum + grad * grad
+// linear += grad + (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
+// quadratic = 1.0 / (accum_new^(lr_power) * lr) + 2 * l2
+// var = (sign(linear) * l1 - linear) / quadratic if |linear| > l1 else 0.0
+// accum = accum_new
+//
+// Arguments:
+// var_: Should be from a Variable().
+// accum: Should be from a Variable().
+// linear: Should be from a Variable().
+// grad: The gradient.
+// indices: A vector of indices into the first dimension of var and accum.
+// lr: Scaling factor. Must be a scalar.
+// l1: L1 regularization. Must be a scalar.
+// l2: L2 regularization. Must be a scalar.
+// lr_power: Scaling factor. Must be a scalar.
+//
+// Returns the created operation.
+func ResourceSparseApplyFtrl(scope *Scope, var_ tf.Output, accum tf.Output, linear tf.Output, grad tf.Output, indices tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, lr_power tf.Output, optional ...ResourceSparseApplyFtrlAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceSparseApplyFtrl",
+ Input: []tf.Input{
+ var_, accum, linear, grad, indices, lr, l1, l2, lr_power,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Returns which elements of x are Inf.
+//
+// @compatibility(numpy)
+// Equivalent to np.isinf
+// @end_compatibility
+func IsInf(scope *Scope, x tf.Output) (y tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "IsInf",
+ Input: []tf.Input{
+ x,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// TruncatedNormalAttr is an optional argument to TruncatedNormal.
+type TruncatedNormalAttr func(optionalAttr)
+
+// TruncatedNormalSeed sets the optional seed attribute to value.
+//
+// value: If either `seed` or `seed2` are set to be non-zero, the random number
+// generator is seeded by the given seed. Otherwise, it is seeded by a
+// random seed.
+// If not specified, defaults to 0
+func TruncatedNormalSeed(value int64) TruncatedNormalAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// TruncatedNormalSeed2 sets the optional seed2 attribute to value.
+//
+// value: A second seed to avoid seed collision.
+// If not specified, defaults to 0
+func TruncatedNormalSeed2(value int64) TruncatedNormalAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// Outputs random values from a truncated normal distribution.
+//
+// The generated values follow a normal distribution with mean 0 and standard
+// deviation 1, except that values whose magnitude is more than 2 standard
+// deviations from the mean are dropped and re-picked.
+//
+// Arguments:
+// shape: The shape of the output tensor.
+// dtype: The type of the output.
+//
+// Returns A tensor of the specified shape filled with random truncated normal
+// values.
+func TruncatedNormal(scope *Scope, shape tf.Output, dtype tf.DataType, optional ...TruncatedNormalAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtype": dtype}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "TruncatedNormal",
+ Input: []tf.Input{
+ shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// SkipgramAttr is an optional argument to Skipgram.
+type SkipgramAttr func(optionalAttr)
+
+// SkipgramWindowSize sets the optional window_size attribute to value.
+//
+// value: The number of words to predict to the left and right of the target.
+// If not specified, defaults to 5
+func SkipgramWindowSize(value int64) SkipgramAttr {
+ return func(m optionalAttr) {
+ m["window_size"] = value
+ }
+}
+
+// SkipgramMinCount sets the optional min_count attribute to value.
+//
+// value: The minimum number of word occurrences for it to be included in the
+// vocabulary.
+// If not specified, defaults to 5
+func SkipgramMinCount(value int64) SkipgramAttr {
+ return func(m optionalAttr) {
+ m["min_count"] = value
+ }
+}
+
+// SkipgramSubsample sets the optional subsample attribute to value.
+//
+// value: Threshold for word occurrence. Words that appear with higher
+// frequency will be randomly down-sampled. Set to 0 to disable.
+// If not specified, defaults to 0.001
+func SkipgramSubsample(value float32) SkipgramAttr {
+ return func(m optionalAttr) {
+ m["subsample"] = value
+ }
+}
+
+// Parses a text file and creates a batch of examples.
+//
+// DEPRECATED at GraphDef version 19: Moving word2vec into tensorflow_models/tutorials and deprecating its ops here as a result
+//
+// Arguments:
+// filename: The corpus's text file name.
+// batch_size: The size of produced batch.
+//
+// Returns A vector of words in the corpus.Frequencies of words. Sorted in the non-ascending order.Number of words per epoch in the data file.The current epoch number.The total number of words processed so far.A vector of word ids.A vector of word ids.
+func Skipgram(scope *Scope, filename string, batch_size int64, optional ...SkipgramAttr) (vocab_word tf.Output, vocab_freq tf.Output, words_per_epoch tf.Output, current_epoch tf.Output, total_words_processed tf.Output, examples tf.Output, labels tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"filename": filename, "batch_size": batch_size}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "Skipgram",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2), op.Output(3), op.Output(4), op.Output(5), op.Output(6)
+}
+
+// StringToNumberAttr is an optional argument to StringToNumber.
+type StringToNumberAttr func(optionalAttr)
+
+// StringToNumberOutType sets the optional out_type attribute to value.
+//
+// value: The numeric type to interpret each string in `string_tensor` as.
+// If not specified, defaults to DT_FLOAT
+func StringToNumberOutType(value tf.DataType) StringToNumberAttr {
+ return func(m optionalAttr) {
+ m["out_type"] = value
+ }
+}
+
+// Converts each string in the input Tensor to the specified numeric type.
+//
+// (Note that int32 overflow results in an error while float overflow
+// results in a rounded value.)
+//
+// Returns A Tensor of the same shape as the input `string_tensor`.
+func StringToNumber(scope *Scope, string_tensor tf.Output, optional ...StringToNumberAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "StringToNumber",
+ Input: []tf.Input{
+ string_tensor,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ResourceApplyFtrlV2Attr is an optional argument to ResourceApplyFtrlV2.
+type ResourceApplyFtrlV2Attr func(optionalAttr)
+
+// ResourceApplyFtrlV2UseLocking sets the optional use_locking attribute to value.
+//
+// value: If `True`, updating of the var and accum tensors will be protected
+// by a lock; otherwise the behavior is undefined, but may exhibit less
+// contention.
+// If not specified, defaults to false
+func ResourceApplyFtrlV2UseLocking(value bool) ResourceApplyFtrlV2Attr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Update '*var' according to the Ftrl-proximal scheme.
+//
+// grad_with_shrinkage = grad + 2 * l2_shrinkage * var
+// accum_new = accum + grad_with_shrinkage * grad_with_shrinkage
+// linear += grad_with_shrinkage +
+// (accum_new^(-lr_power) - accum^(-lr_power)) / lr * var
+// quadratic = 1.0 / (accum_new^(lr_power) * lr) + 2 * l2
+// var = (sign(linear) * l1 - linear) / quadratic if |linear| > l1 else 0.0
+// accum = accum_new
+//
+// Arguments:
+// var_: Should be from a Variable().
+// accum: Should be from a Variable().
+// linear: Should be from a Variable().
+// grad: The gradient.
+// lr: Scaling factor. Must be a scalar.
+// l1: L1 regulariation. Must be a scalar.
+// l2: L2 shrinkage regulariation. Must be a scalar.
+//
+// lr_power: Scaling factor. Must be a scalar.
+//
+// Returns the created operation.
+func ResourceApplyFtrlV2(scope *Scope, var_ tf.Output, accum tf.Output, linear tf.Output, grad tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, l2_shrinkage tf.Output, lr_power tf.Output, optional ...ResourceApplyFtrlV2Attr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceApplyFtrlV2",
+ Input: []tf.Input{
+ var_, accum, linear, grad, lr, l1, l2, l2_shrinkage, lr_power,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// EncodeJpegAttr is an optional argument to EncodeJpeg.
+type EncodeJpegAttr func(optionalAttr)
+
+// EncodeJpegFormat sets the optional format attribute to value.
+//
+// value: Per pixel image format.
+// If not specified, defaults to ""
+func EncodeJpegFormat(value string) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["format"] = value
+ }
+}
+
+// EncodeJpegQuality sets the optional quality attribute to value.
+//
+// value: Quality of the compression from 0 to 100 (higher is better and slower).
+// If not specified, defaults to 95
+func EncodeJpegQuality(value int64) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["quality"] = value
+ }
+}
+
+// EncodeJpegProgressive sets the optional progressive attribute to value.
+//
+// value: If True, create a JPEG that loads progressively (coarse to fine).
+// If not specified, defaults to false
+func EncodeJpegProgressive(value bool) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["progressive"] = value
+ }
+}
+
+// EncodeJpegOptimizeSize sets the optional optimize_size attribute to value.
+//
+// value: If True, spend CPU/RAM to reduce size with no quality change.
+// If not specified, defaults to false
+func EncodeJpegOptimizeSize(value bool) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["optimize_size"] = value
+ }
+}
+
+// EncodeJpegChromaDownsampling sets the optional chroma_downsampling attribute to value.
+//
+// value: See http://en.wikipedia.org/wiki/Chroma_subsampling.
+// If not specified, defaults to true
+func EncodeJpegChromaDownsampling(value bool) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["chroma_downsampling"] = value
+ }
+}
+
+// EncodeJpegDensityUnit sets the optional density_unit attribute to value.
+//
+// value: Unit used to specify `x_density` and `y_density`:
+// pixels per inch (`'in'`) or centimeter (`'cm'`).
+// If not specified, defaults to "in"
+func EncodeJpegDensityUnit(value string) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["density_unit"] = value
+ }
+}
+
+// EncodeJpegXDensity sets the optional x_density attribute to value.
+//
+// value: Horizontal pixels per density unit.
+// If not specified, defaults to 300
+func EncodeJpegXDensity(value int64) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["x_density"] = value
+ }
+}
+
+// EncodeJpegYDensity sets the optional y_density attribute to value.
+//
+// value: Vertical pixels per density unit.
+// If not specified, defaults to 300
+func EncodeJpegYDensity(value int64) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["y_density"] = value
+ }
+}
+
+// EncodeJpegXmpMetadata sets the optional xmp_metadata attribute to value.
+//
+// value: If not empty, embed this XMP metadata in the image header.
+// If not specified, defaults to ""
+func EncodeJpegXmpMetadata(value string) EncodeJpegAttr {
+ return func(m optionalAttr) {
+ m["xmp_metadata"] = value
+ }
+}
+
+// JPEG-encode an image.
+//
+// `image` is a 3-D uint8 Tensor of shape `[height, width, channels]`.
+//
+// The attr `format` can be used to override the color format of the encoded
+// output. Values can be:
+//
+// * `''`: Use a default format based on the number of channels in the image.
+// * `grayscale`: Output a grayscale JPEG image. The `channels` dimension
+// of `image` must be 1.
+// * `rgb`: Output an RGB JPEG image. The `channels` dimension
+// of `image` must be 3.
+//
+// If `format` is not specified or is the empty string, a default format is picked
+// in function of the number of channels in `image`:
+//
+// * 1: Output a grayscale image.
+// * 3: Output an RGB image.
+//
+// Arguments:
+// image: 3-D with shape `[height, width, channels]`.
+//
+// Returns 0-D. JPEG-encoded image.
+func EncodeJpeg(scope *Scope, image tf.Output, optional ...EncodeJpegAttr) (contents tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "EncodeJpeg",
+ Input: []tf.Input{
+ image,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// MultinomialAttr is an optional argument to Multinomial.
+type MultinomialAttr func(optionalAttr)
+
+// MultinomialSeed sets the optional seed attribute to value.
+//
+// value: If either seed or seed2 is set to be non-zero, the internal random number
+// generator is seeded by the given seed. Otherwise, a random seed is used.
+// If not specified, defaults to 0
+func MultinomialSeed(value int64) MultinomialAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// MultinomialSeed2 sets the optional seed2 attribute to value.
+//
+// value: A second seed to avoid seed collision.
+// If not specified, defaults to 0
+func MultinomialSeed2(value int64) MultinomialAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// MultinomialOutputDtype sets the optional output_dtype attribute to value.
+// If not specified, defaults to DT_INT64
+func MultinomialOutputDtype(value tf.DataType) MultinomialAttr {
+ return func(m optionalAttr) {
+ m["output_dtype"] = value
+ }
+}
+
+// Draws samples from a multinomial distribution.
+//
+// Arguments:
+// logits: 2-D Tensor with shape `[batch_size, num_classes]`. Each slice `[i, :]`
+// represents the unnormalized log probabilities for all classes.
+// num_samples: 0-D. Number of independent samples to draw for each row slice.
+//
+// Returns 2-D Tensor with shape `[batch_size, num_samples]`. Each slice `[i, :]`
+// contains the drawn class labels with range `[0, num_classes)`.
+func Multinomial(scope *Scope, logits tf.Output, num_samples tf.Output, optional ...MultinomialAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "Multinomial",
+ Input: []tf.Input{
+ logits, num_samples,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ResourceSparseApplyAdagradDAAttr is an optional argument to ResourceSparseApplyAdagradDA.
+type ResourceSparseApplyAdagradDAAttr func(optionalAttr)
+
+// ResourceSparseApplyAdagradDAUseLocking sets the optional use_locking attribute to value.
+//
+// value: If True, updating of the var and accum tensors will be protected by
+// a lock; otherwise the behavior is undefined, but may exhibit less contention.
+// If not specified, defaults to false
+func ResourceSparseApplyAdagradDAUseLocking(value bool) ResourceSparseApplyAdagradDAAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Update entries in '*var' and '*accum' according to the proximal adagrad scheme.
+//
+// Arguments:
+// var_: Should be from a Variable().
+// gradient_accumulator: Should be from a Variable().
+// gradient_squared_accumulator: Should be from a Variable().
+// grad: The gradient.
+// indices: A vector of indices into the first dimension of var and accum.
+// lr: Learning rate. Must be a scalar.
+// l1: L1 regularization. Must be a scalar.
+// l2: L2 regularization. Must be a scalar.
+// global_step: Training step number. Must be a scalar.
+//
+// Returns the created operation.
+func ResourceSparseApplyAdagradDA(scope *Scope, var_ tf.Output, gradient_accumulator tf.Output, gradient_squared_accumulator tf.Output, grad tf.Output, indices tf.Output, lr tf.Output, l1 tf.Output, l2 tf.Output, global_step tf.Output, optional ...ResourceSparseApplyAdagradDAAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceSparseApplyAdagradDA",
+ Input: []tf.Input{
+ var_, gradient_accumulator, gradient_squared_accumulator, grad, indices, lr, l1, l2, global_step,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
// Converts each string in the input Tensor to its hash mod by a number of buckets.
//
// The hash function is deterministic on the content of the string within the
@@ -17021,6 +16841,361 @@
return scope.AddOperation(opspec)
}
+// SerializeManySparseAttr is an optional argument to SerializeManySparse.
+type SerializeManySparseAttr func(optionalAttr)
+
+// SerializeManySparseOutType sets the optional out_type attribute to value.
+//
+// value: The `dtype` to use for serialization; the supported types are `string`
+// (default) and `variant`.
+// If not specified, defaults to DT_STRING
+func SerializeManySparseOutType(value tf.DataType) SerializeManySparseAttr {
+ return func(m optionalAttr) {
+ m["out_type"] = value
+ }
+}
+
+// Serialize an `N`-minibatch `SparseTensor` into an `[N, 3]` `Tensor` object.
+//
+// The `SparseTensor` must have rank `R` greater than 1, and the first dimension
+// is treated as the minibatch dimension. Elements of the `SparseTensor`
+// must be sorted in increasing order of this first dimension. The serialized
+// `SparseTensor` objects going into each row of `serialized_sparse` will have
+// rank `R-1`.
+//
+// The minibatch size `N` is extracted from `sparse_shape[0]`.
+//
+// Arguments:
+// sparse_indices: 2-D. The `indices` of the minibatch `SparseTensor`.
+// sparse_values: 1-D. The `values` of the minibatch `SparseTensor`.
+// sparse_shape: 1-D. The `shape` of the minibatch `SparseTensor`.
+func SerializeManySparse(scope *Scope, sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output, optional ...SerializeManySparseAttr) (serialized_sparse tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "SerializeManySparse",
+ Input: []tf.Input{
+ sparse_indices, sparse_values, sparse_shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Computes inverse hyperbolic cosine of x element-wise.
+func Acosh(scope *Scope, x tf.Output) (y tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Acosh",
+ Input: []tf.Input{
+ x,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// TensorArrayV2Attr is an optional argument to TensorArrayV2.
+type TensorArrayV2Attr func(optionalAttr)
+
+// TensorArrayV2ElementShape sets the optional element_shape attribute to value.
+// If not specified, defaults to <unknown_rank:true >
+func TensorArrayV2ElementShape(value tf.Shape) TensorArrayV2Attr {
+ return func(m optionalAttr) {
+ m["element_shape"] = value
+ }
+}
+
+// TensorArrayV2DynamicSize sets the optional dynamic_size attribute to value.
+// If not specified, defaults to false
+func TensorArrayV2DynamicSize(value bool) TensorArrayV2Attr {
+ return func(m optionalAttr) {
+ m["dynamic_size"] = value
+ }
+}
+
+// TensorArrayV2ClearAfterRead sets the optional clear_after_read attribute to value.
+// If not specified, defaults to true
+func TensorArrayV2ClearAfterRead(value bool) TensorArrayV2Attr {
+ return func(m optionalAttr) {
+ m["clear_after_read"] = value
+ }
+}
+
+// TensorArrayV2TensorArrayName sets the optional tensor_array_name attribute to value.
+// If not specified, defaults to ""
+func TensorArrayV2TensorArrayName(value string) TensorArrayV2Attr {
+ return func(m optionalAttr) {
+ m["tensor_array_name"] = value
+ }
+}
+
+// Deprecated. Use TensorArrayV3
+//
+// DEPRECATED at GraphDef version 26: Use TensorArrayV3
+func TensorArrayV2(scope *Scope, size tf.Output, dtype tf.DataType, optional ...TensorArrayV2Attr) (handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtype": dtype}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "TensorArrayV2",
+ Input: []tf.Input{
+ size,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ThreadUnsafeUnigramCandidateSamplerAttr is an optional argument to ThreadUnsafeUnigramCandidateSampler.
+type ThreadUnsafeUnigramCandidateSamplerAttr func(optionalAttr)
+
+// ThreadUnsafeUnigramCandidateSamplerSeed sets the optional seed attribute to value.
+//
+// value: If either seed or seed2 are set to be non-zero, the random number
+// generator is seeded by the given seed. Otherwise, it is seeded by a
+// random seed.
+// If not specified, defaults to 0
+func ThreadUnsafeUnigramCandidateSamplerSeed(value int64) ThreadUnsafeUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// ThreadUnsafeUnigramCandidateSamplerSeed2 sets the optional seed2 attribute to value.
+//
+// value: An second seed to avoid seed collision.
+// If not specified, defaults to 0
+func ThreadUnsafeUnigramCandidateSamplerSeed2(value int64) ThreadUnsafeUnigramCandidateSamplerAttr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// Generates labels for candidate sampling with a learned unigram distribution.
+//
+// See explanations of candidate sampling and the data formats at
+// go/candidate-sampling.
+//
+// For each batch, this op picks a single set of sampled candidate labels.
+//
+// The advantages of sampling candidates per-batch are simplicity and the
+// possibility of efficient dense matrix multiplication. The disadvantage is that
+// the sampled candidates must be chosen independently of the context and of the
+// true labels.
+//
+// Arguments:
+// true_classes: A batch_size * num_true matrix, in which each row contains the
+// IDs of the num_true target_classes in the corresponding original label.
+// num_true: Number of true labels per context.
+// num_sampled: Number of candidates to randomly sample.
+// unique: If unique is true, we sample with rejection, so that all sampled
+// candidates in a batch are unique. This requires some approximation to
+// estimate the post-rejection sampling probabilities.
+// range_max: The sampler will sample integers from the interval [0, range_max).
+//
+// Returns A vector of length num_sampled, in which each element is
+// the ID of a sampled candidate.A batch_size * num_true matrix, representing
+// the number of times each candidate is expected to occur in a batch
+// of sampled candidates. If unique=true, then this is a probability.A vector of length num_sampled, for each sampled
+// candidate representing the number of times the candidate is expected
+// to occur in a batch of sampled candidates. If unique=true, then this is a
+// probability.
+func ThreadUnsafeUnigramCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...ThreadUnsafeUnigramCandidateSamplerAttr) (sampled_candidates tf.Output, true_expected_count tf.Output, sampled_expected_count tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"num_true": num_true, "num_sampled": num_sampled, "unique": unique, "range_max": range_max}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ThreadUnsafeUnigramCandidateSampler",
+ Input: []tf.Input{
+ true_classes,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// MaxPoolV2Attr is an optional argument to MaxPoolV2.
+type MaxPoolV2Attr func(optionalAttr)
+
+// MaxPoolV2DataFormat sets the optional data_format attribute to value.
+//
+// value: Specify the data format of the input and output data. With the
+// default format "NHWC", the data is stored in the order of:
+// [batch, in_height, in_width, in_channels].
+// Alternatively, the format could be "NCHW", the data storage order of:
+// [batch, in_channels, in_height, in_width].
+// If not specified, defaults to "NHWC"
+func MaxPoolV2DataFormat(value string) MaxPoolV2Attr {
+ return func(m optionalAttr) {
+ m["data_format"] = value
+ }
+}
+
+// Performs max pooling on the input.
+//
+// Arguments:
+// input: 4-D input to pool over.
+// ksize: The size of the window for each dimension of the input tensor.
+// strides: The stride of the sliding window for each dimension of the
+// input tensor.
+// padding: The type of padding algorithm to use.
+//
+// Returns The max pooled output tensor.
+func MaxPoolV2(scope *Scope, input tf.Output, ksize tf.Output, strides tf.Output, padding string, optional ...MaxPoolV2Attr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"padding": padding}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MaxPoolV2",
+ Input: []tf.Input{
+ input, ksize, strides,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// AddManySparseToTensorsMapAttr is an optional argument to AddManySparseToTensorsMap.
+type AddManySparseToTensorsMapAttr func(optionalAttr)
+
+// AddManySparseToTensorsMapContainer sets the optional container attribute to value.
+//
+// value: The container name for the `SparseTensorsMap` created by this op.
+// If not specified, defaults to ""
+func AddManySparseToTensorsMapContainer(value string) AddManySparseToTensorsMapAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// AddManySparseToTensorsMapSharedName sets the optional shared_name attribute to value.
+//
+// value: The shared name for the `SparseTensorsMap` created by this op.
+// If blank, the new Operation's unique name is used.
+// If not specified, defaults to ""
+func AddManySparseToTensorsMapSharedName(value string) AddManySparseToTensorsMapAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Add an `N`-minibatch `SparseTensor` to a `SparseTensorsMap`, return `N` handles.
+//
+// A `SparseTensor` of rank `R` is represented by three tensors: `sparse_indices`,
+// `sparse_values`, and `sparse_shape`, where
+//
+// ```sparse_indices.shape[1] == sparse_shape.shape[0] == R```
+//
+// An `N`-minibatch of `SparseTensor` objects is represented as a `SparseTensor`
+// having a first `sparse_indices` column taking values between `[0, N)`, where
+// the minibatch size `N == sparse_shape[0]`.
+//
+// The input `SparseTensor` must have rank `R` greater than 1, and the first
+// dimension is treated as the minibatch dimension. Elements of the `SparseTensor`
+// must be sorted in increasing order of this first dimension. The stored
+// `SparseTensor` objects pointed to by each row of the output `sparse_handles`
+// will have rank `R-1`.
+//
+// The `SparseTensor` values can then be read out as part of a minibatch by passing
+// the given keys as vector elements to `TakeManySparseFromTensorsMap`. To ensure
+// the correct `SparseTensorsMap` is accessed, ensure that the same
+// `container` and `shared_name` are passed to that Op. If no `shared_name`
+// is provided here, instead use the *name* of the Operation created by calling
+// `AddManySparseToTensorsMap` as the `shared_name` passed to
+// `TakeManySparseFromTensorsMap`. Ensure the Operations are colocated.
+//
+// Arguments:
+// sparse_indices: 2-D. The `indices` of the minibatch `SparseTensor`.
+// `sparse_indices[:, 0]` must be ordered values in `[0, N)`.
+// sparse_values: 1-D. The `values` of the minibatch `SparseTensor`.
+// sparse_shape: 1-D. The `shape` of the minibatch `SparseTensor`.
+// The minibatch size `N == sparse_shape[0]`.
+//
+// Returns 1-D. The handles of the `SparseTensor` now stored in the
+// `SparseTensorsMap`. Shape: `[N]`.
+func AddManySparseToTensorsMap(scope *Scope, sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output, optional ...AddManySparseToTensorsMapAttr) (sparse_handles tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "AddManySparseToTensorsMap",
+ Input: []tf.Input{
+ sparse_indices, sparse_values, sparse_shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Concatenates tensors along one dimension.
+//
+// Arguments:
+// values: List of `N` Tensors to concatenate. Their ranks and types must match,
+// and their sizes must match in all dimensions except `concat_dim`.
+// axis: 0-D. The dimension along which to concatenate. Must be in the
+// range [-rank(values), rank(values)).
+//
+// Returns A `Tensor` with the concatenation of values stacked along the
+// `concat_dim` dimension. This tensor's shape matches that of `values` except
+// in `concat_dim` where it has the sum of the sizes.
+func ConcatV2(scope *Scope, values []tf.Output, axis tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ConcatV2",
+ Input: []tf.Input{
+ tf.OutputList(values), axis,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Reads and outputs the entire contents of the input filename.
+func ReadFile(scope *Scope, filename tf.Output) (contents tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "ReadFile",
+ Input: []tf.Input{
+ filename,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Returns immutable tensor from memory region.
//
// The current implementation memmaps the tensor from a file.
@@ -17108,6 +17283,44 @@
return op.Output(0)
}
+// Computes softsign gradients for a softsign operation.
+//
+// Arguments:
+// gradients: The backpropagated gradients to the corresponding softsign operation.
+// features: The features passed as input to the corresponding softsign operation.
+//
+// Returns The gradients: `gradients / (1 + abs(features)) ** 2`.
+func SoftsignGrad(scope *Scope, gradients tf.Output, features tf.Output) (backprops tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SoftsignGrad",
+ Input: []tf.Input{
+ gradients, features,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Provides the time since epoch in seconds.
+//
+// Returns the timestamp as a `float64` for seconds since the Unix epoch.
+//
+// Note: the timestamp is computed when the op is executed, not when it is added
+// to the graph.
+func Timestamp(scope *Scope) (ts tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Timestamp",
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// VariableShapeAttr is an optional argument to VariableShape.
type VariableShapeAttr func(optionalAttr)
@@ -17148,80 +17361,96 @@
return op.Output(0)
}
-// SparseToSparseSetOperationAttr is an optional argument to SparseToSparseSetOperation.
-type SparseToSparseSetOperationAttr func(optionalAttr)
+// AvgPoolGradAttr is an optional argument to AvgPoolGrad.
+type AvgPoolGradAttr func(optionalAttr)
-// SparseToSparseSetOperationValidateIndices sets the optional validate_indices attribute to value.
-// If not specified, defaults to true
-func SparseToSparseSetOperationValidateIndices(value bool) SparseToSparseSetOperationAttr {
+// AvgPoolGradDataFormat sets the optional data_format attribute to value.
+//
+// value: Specify the data format of the input and output data. With the
+// default format "NHWC", the data is stored in the order of:
+// [batch, in_height, in_width, in_channels].
+// Alternatively, the format could be "NCHW", the data storage order of:
+// [batch, in_channels, in_height, in_width].
+// If not specified, defaults to "NHWC"
+func AvgPoolGradDataFormat(value string) AvgPoolGradAttr {
return func(m optionalAttr) {
- m["validate_indices"] = value
+ m["data_format"] = value
}
}
-// Applies set operation along last dimension of 2 `SparseTensor` inputs.
-//
-// See SetOperationOp::SetOperationFromContext for values of `set_operation`.
-//
-// If `validate_indices` is `True`, `SparseToSparseSetOperation` validates the
-// order and range of `set1` and `set2` indices.
-//
-// Input `set1` is a `SparseTensor` represented by `set1_indices`, `set1_values`,
-// and `set1_shape`. For `set1` ranked `n`, 1st `n-1` dimensions must be the same
-// as `set2`. Dimension `n` contains values in a set, duplicates are allowed but
-// ignored.
-//
-// Input `set2` is a `SparseTensor` represented by `set2_indices`, `set2_values`,
-// and `set2_shape`. For `set2` ranked `n`, 1st `n-1` dimensions must be the same
-// as `set1`. Dimension `n` contains values in a set, duplicates are allowed but
-// ignored.
-//
-// If `validate_indices` is `True`, this op validates the order and range of `set1`
-// and `set2` indices.
-//
-// Output `result` is a `SparseTensor` represented by `result_indices`,
-// `result_values`, and `result_shape`. For `set1` and `set2` ranked `n`, this
-// has rank `n` and the same 1st `n-1` dimensions as `set1` and `set2`. The `nth`
-// dimension contains the result of `set_operation` applied to the corresponding
-// `[0...n-1]` dimension of `set`.
+// Computes gradients of the average pooling function.
//
// Arguments:
-// set1_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major
-// order.
-// set1_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major
-// order.
-// set1_shape: 1D `Tensor`, shape of a `SparseTensor`. `set1_shape[0...n-1]` must
-// be the same as `set2_shape[0...n-1]`, `set1_shape[n]` is the
-// max set size across `0...n-1` dimensions.
-// set2_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major
-// order.
-// set2_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major
-// order.
-// set2_shape: 1D `Tensor`, shape of a `SparseTensor`. `set2_shape[0...n-1]` must
-// be the same as `set1_shape[0...n-1]`, `set2_shape[n]` is the
-// max set size across `0...n-1` dimensions.
+// orig_input_shape: 1-D. Shape of the original input to `avg_pool`.
+// grad: 4-D with shape `[batch, height, width, channels]`. Gradients w.r.t.
+// the output of `avg_pool`.
+// ksize: The size of the sliding window for each dimension of the input.
+// strides: The stride of the sliding window for each dimension of the input.
+// padding: The type of padding algorithm to use.
//
-//
-// Returns 2D indices of a `SparseTensor`.1D values of a `SparseTensor`.1D `Tensor` shape of a `SparseTensor`. `result_shape[0...n-1]` is
-// the same as the 1st `n-1` dimensions of `set1` and `set2`, `result_shape[n]`
-// is the max result set size across all `0...n-1` dimensions.
-func SparseToSparseSetOperation(scope *Scope, set1_indices tf.Output, set1_values tf.Output, set1_shape tf.Output, set2_indices tf.Output, set2_values tf.Output, set2_shape tf.Output, set_operation string, optional ...SparseToSparseSetOperationAttr) (result_indices tf.Output, result_values tf.Output, result_shape tf.Output) {
+// Returns 4-D. Gradients w.r.t. the input of `avg_pool`.
+func AvgPoolGrad(scope *Scope, orig_input_shape tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...AvgPoolGradAttr) (output tf.Output) {
if scope.Err() != nil {
return
}
- attrs := map[string]interface{}{"set_operation": set_operation}
+ attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding}
for _, a := range optional {
a(attrs)
}
opspec := tf.OpSpec{
- Type: "SparseToSparseSetOperation",
+ Type: "AvgPoolGrad",
Input: []tf.Input{
- set1_indices, set1_values, set1_shape, set2_indices, set2_values, set2_shape,
+ orig_input_shape, grad,
},
Attrs: attrs,
}
op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
+ return op.Output(0)
+}
+
+// Greedily selects a subset of bounding boxes in descending order of score,
+//
+// pruning away boxes that have high overlaps
+// with previously selected boxes. Bounding boxes with score less than
+// `score_threshold` are removed. N-by-n overlap values are supplied as square matrix,
+// which allows for defining a custom overlap criterium (eg. intersection over union,
+// intersection over area, etc.).
+//
+// The output of this operation is a set of integers indexing into the input
+// collection of bounding boxes representing the selected boxes. The bounding
+// box coordinates corresponding to the selected indices can then be obtained
+// using the `tf.gather operation`. For example:
+//
+// selected_indices = tf.image.non_max_suppression_with_overlaps(
+// overlaps, scores, max_output_size, overlap_threshold, score_threshold)
+// selected_boxes = tf.gather(boxes, selected_indices)
+//
+// Arguments:
+// overlaps: A 2-D float tensor of shape `[num_boxes, num_boxes]` representing
+// the n-by-n box overlap values.
+// scores: A 1-D float tensor of shape `[num_boxes]` representing a single
+// score corresponding to each box (each row of boxes).
+// max_output_size: A scalar integer tensor representing the maximum number of
+// boxes to be selected by non max suppression.
+// overlap_threshold: A 0-D float tensor representing the threshold for deciding whether
+// boxes overlap too.
+// score_threshold: A 0-D float tensor representing the threshold for deciding when to remove
+// boxes based on score.
+//
+// Returns A 1-D integer tensor of shape `[M]` representing the selected
+// indices from the boxes tensor, where `M <= max_output_size`.
+func NonMaxSuppressionWithOverlaps(scope *Scope, overlaps tf.Output, scores tf.Output, max_output_size tf.Output, overlap_threshold tf.Output, score_threshold tf.Output) (selected_indices tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "NonMaxSuppressionWithOverlaps",
+ Input: []tf.Input{
+ overlaps, scores, max_output_size, overlap_threshold, score_threshold,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
}
// Computes softmax cross entropy cost and gradients to backpropagate.
@@ -17615,47 +17844,6 @@
return op.Output(0)
}
-// Subtracts sparse updates from the variable referenced by `resource`.
-//
-// This operation computes
-//
-// # Scalar indices
-// ref[indices, ...] -= updates[...]
-//
-// # Vector indices (for each i)
-// ref[indices[i], ...] -= updates[i, ...]
-//
-// # High rank indices (for each i, ..., j)
-// ref[indices[i, ..., j], ...] -= updates[i, ..., j, ...]
-//
-// Duplicate entries are handled correctly: if multiple `indices` reference
-// the same location, their contributions add.
-//
-// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`.
-//
-// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
-// <img style="width:100%" src='https://www.tensorflow.org/images/ScatterAdd.png' alt>
-// </div>
-//
-// Arguments:
-// resource: Should be from a `Variable` node.
-// indices: A tensor of indices into the first dimension of `ref`.
-// updates: A tensor of updated values to add to `ref`.
-//
-// Returns the created operation.
-func ResourceScatterSub(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ResourceScatterSub",
- Input: []tf.Input{
- resource, indices, updates,
- },
- }
- return scope.AddOperation(opspec)
-}
-
// DepthwiseConv2dNativeBackpropFilterAttr is an optional argument to DepthwiseConv2dNativeBackpropFilter.
type DepthwiseConv2dNativeBackpropFilterAttr func(optionalAttr)
@@ -18090,6 +18278,93 @@
return op.Output(0)
}
+// DenseToSparseSetOperationAttr is an optional argument to DenseToSparseSetOperation.
+type DenseToSparseSetOperationAttr func(optionalAttr)
+
+// DenseToSparseSetOperationValidateIndices sets the optional validate_indices attribute to value.
+// If not specified, defaults to true
+func DenseToSparseSetOperationValidateIndices(value bool) DenseToSparseSetOperationAttr {
+ return func(m optionalAttr) {
+ m["validate_indices"] = value
+ }
+}
+
+// Applies set operation along last dimension of `Tensor` and `SparseTensor`.
+//
+// See SetOperationOp::SetOperationFromContext for values of `set_operation`.
+//
+// Input `set2` is a `SparseTensor` represented by `set2_indices`, `set2_values`,
+// and `set2_shape`. For `set2` ranked `n`, 1st `n-1` dimensions must be the same
+// as `set1`. Dimension `n` contains values in a set, duplicates are allowed but
+// ignored.
+//
+// If `validate_indices` is `True`, this op validates the order and range of `set2`
+// indices.
+//
+// Output `result` is a `SparseTensor` represented by `result_indices`,
+// `result_values`, and `result_shape`. For `set1` and `set2` ranked `n`, this
+// has rank `n` and the same 1st `n-1` dimensions as `set1` and `set2`. The `nth`
+// dimension contains the result of `set_operation` applied to the corresponding
+// `[0...n-1]` dimension of `set`.
+//
+// Arguments:
+// set1: `Tensor` with rank `n`. 1st `n-1` dimensions must be the same as `set2`.
+// Dimension `n` contains values in a set, duplicates are allowed but ignored.
+// set2_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major
+// order.
+// set2_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major
+// order.
+// set2_shape: 1D `Tensor`, shape of a `SparseTensor`. `set2_shape[0...n-1]` must
+// be the same as the 1st `n-1` dimensions of `set1`, `result_shape[n]` is the
+// max set size across `n-1` dimensions.
+//
+//
+// Returns 2D indices of a `SparseTensor`.1D values of a `SparseTensor`.1D `Tensor` shape of a `SparseTensor`. `result_shape[0...n-1]` is
+// the same as the 1st `n-1` dimensions of `set1` and `set2`, `result_shape[n]`
+// is the max result set size across all `0...n-1` dimensions.
+func DenseToSparseSetOperation(scope *Scope, set1 tf.Output, set2_indices tf.Output, set2_values tf.Output, set2_shape tf.Output, set_operation string, optional ...DenseToSparseSetOperationAttr) (result_indices tf.Output, result_values tf.Output, result_shape tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"set_operation": set_operation}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "DenseToSparseSetOperation",
+ Input: []tf.Input{
+ set1, set2_indices, set2_values, set2_shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// L2 Loss.
+//
+// Computes half the L2 norm of a tensor without the `sqrt`:
+//
+// output = sum(t ** 2) / 2
+//
+// Arguments:
+// t: Typically 2-D, but may have any dimensions.
+//
+// Returns 0-D.
+func L2Loss(scope *Scope, t tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "L2Loss",
+ Input: []tf.Input{
+ t,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Computes cos of x element-wise.
func Cos(scope *Scope, x tf.Output) (y tf.Output) {
if scope.Err() != nil {
@@ -18236,6 +18511,38 @@
return op.Output(0), op.Output(1)
}
+// The gradient operator for the SparseAdd op.
+//
+// The SparseAdd op calculates A + B, where A, B, and the sum are all represented
+// as `SparseTensor` objects. This op takes in the upstream gradient w.r.t.
+// non-empty values of the sum, and outputs the gradients w.r.t. the non-empty
+// values of A and B.
+//
+// Arguments:
+// backprop_val_grad: 1-D with shape `[nnz(sum)]`. The gradient with respect to
+// the non-empty values of the sum.
+// a_indices: 2-D. The `indices` of the `SparseTensor` A, size `[nnz(A), ndims]`.
+// b_indices: 2-D. The `indices` of the `SparseTensor` B, size `[nnz(B), ndims]`.
+// sum_indices: 2-D. The `indices` of the sum `SparseTensor`, size
+// `[nnz(sum), ndims]`.
+//
+// Returns 1-D with shape `[nnz(A)]`. The gradient with respect to the
+// non-empty values of A.1-D with shape `[nnz(B)]`. The gradient with respect to the
+// non-empty values of B.
+func SparseAddGrad(scope *Scope, backprop_val_grad tf.Output, a_indices tf.Output, b_indices tf.Output, sum_indices tf.Output) (a_val_grad tf.Output, b_val_grad tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseAddGrad",
+ Input: []tf.Input{
+ backprop_val_grad, a_indices, b_indices, sum_indices,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1)
+}
+
// Returns a list of tensors with the same shapes and contents as the input
//
// tensors.
@@ -18342,6 +18649,451 @@
return scope.AddOperation(opspec)
}
+// ResourceSparseApplyCenteredRMSPropAttr is an optional argument to ResourceSparseApplyCenteredRMSProp.
+type ResourceSparseApplyCenteredRMSPropAttr func(optionalAttr)
+
+// ResourceSparseApplyCenteredRMSPropUseLocking sets the optional use_locking attribute to value.
+//
+// value: If `True`, updating of the var, mg, ms, and mom tensors is
+// protected by a lock; otherwise the behavior is undefined, but may exhibit less
+// contention.
+// If not specified, defaults to false
+func ResourceSparseApplyCenteredRMSPropUseLocking(value bool) ResourceSparseApplyCenteredRMSPropAttr {
+ return func(m optionalAttr) {
+ m["use_locking"] = value
+ }
+}
+
+// Update '*var' according to the centered RMSProp algorithm.
+//
+// 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)
+//
+// ms <- rho * ms_{t-1} + (1-rho) * grad * grad
+// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms + epsilon)
+// var <- var - mom
+//
+// Arguments:
+// var_: Should be from a Variable().
+// mg: Should be from a Variable().
+// ms: Should be from a Variable().
+// mom: Should be from a Variable().
+// lr: Scaling factor. Must be a scalar.
+// rho: Decay rate. Must be a scalar.
+//
+// epsilon: Ridge term. Must be a scalar.
+// grad: The gradient.
+// indices: A vector of indices into the first dimension of var, ms and mom.
+//
+// Returns the created operation.
+func ResourceSparseApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms tf.Output, mom tf.Output, lr tf.Output, rho tf.Output, momentum tf.Output, epsilon tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyCenteredRMSPropAttr) (o *tf.Operation) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ResourceSparseApplyCenteredRMSProp",
+ Input: []tf.Input{
+ var_, mg, ms, mom, lr, rho, momentum, epsilon, grad, indices,
+ },
+ Attrs: attrs,
+ }
+ return scope.AddOperation(opspec)
+}
+
+// Creates a dataset that batches `batch_size` elements from `input_dataset`.
+//
+// Arguments:
+//
+// batch_size: A scalar representing the number of elements to accumulate in a
+// batch.
+//
+//
+func BatchDataset(scope *Scope, input_dataset tf.Output, batch_size tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes}
+ opspec := tf.OpSpec{
+ Type: "BatchDataset",
+ Input: []tf.Input{
+ input_dataset, batch_size,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// RandomPoissonV2Attr is an optional argument to RandomPoissonV2.
+type RandomPoissonV2Attr func(optionalAttr)
+
+// RandomPoissonV2Seed sets the optional seed attribute to value.
+//
+// value: If either `seed` or `seed2` are set to be non-zero, the random number
+// generator is seeded by the given seed. Otherwise, it is seeded by a
+// random seed.
+// If not specified, defaults to 0
+func RandomPoissonV2Seed(value int64) RandomPoissonV2Attr {
+ return func(m optionalAttr) {
+ m["seed"] = value
+ }
+}
+
+// RandomPoissonV2Seed2 sets the optional seed2 attribute to value.
+//
+// value: A second seed to avoid seed collision.
+// If not specified, defaults to 0
+func RandomPoissonV2Seed2(value int64) RandomPoissonV2Attr {
+ return func(m optionalAttr) {
+ m["seed2"] = value
+ }
+}
+
+// RandomPoissonV2Dtype sets the optional dtype attribute to value.
+// If not specified, defaults to DT_INT64
+func RandomPoissonV2Dtype(value tf.DataType) RandomPoissonV2Attr {
+ return func(m optionalAttr) {
+ m["dtype"] = value
+ }
+}
+
+// Outputs random values from the Poisson distribution(s) described by rate.
+//
+// This op uses two algorithms, depending on rate. If rate >= 10, then
+// the algorithm by Hormann is used to acquire samples via
+// transformation-rejection.
+// See http://www.sciencedirect.com/science/article/pii/0167668793909974.
+//
+// Otherwise, Knuth's algorithm is used to acquire samples via multiplying uniform
+// random variables.
+// See Donald E. Knuth (1969). Seminumerical Algorithms. The Art of Computer
+// Programming, Volume 2. Addison Wesley
+//
+// Arguments:
+// shape: 1-D integer tensor. Shape of independent samples to draw from each
+// distribution described by the shape parameters given in rate.
+// rate: A tensor in which each scalar is a "rate" parameter describing the
+// associated poisson distribution.
+//
+// Returns A tensor with shape `shape + shape(rate)`. Each slice
+// `[:, ..., :, i0, i1, ...iN]` contains the samples drawn for
+// `rate[i0, i1, ...iN]`.
+func RandomPoissonV2(scope *Scope, shape tf.Output, rate tf.Output, optional ...RandomPoissonV2Attr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "RandomPoissonV2",
+ Input: []tf.Input{
+ shape, rate,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// DecodeAndCropJpegAttr is an optional argument to DecodeAndCropJpeg.
+type DecodeAndCropJpegAttr func(optionalAttr)
+
+// DecodeAndCropJpegChannels sets the optional channels attribute to value.
+//
+// value: Number of color channels for the decoded image.
+// If not specified, defaults to 0
+func DecodeAndCropJpegChannels(value int64) DecodeAndCropJpegAttr {
+ return func(m optionalAttr) {
+ m["channels"] = value
+ }
+}
+
+// DecodeAndCropJpegRatio sets the optional ratio attribute to value.
+//
+// value: Downscaling ratio.
+// If not specified, defaults to 1
+func DecodeAndCropJpegRatio(value int64) DecodeAndCropJpegAttr {
+ return func(m optionalAttr) {
+ m["ratio"] = value
+ }
+}
+
+// DecodeAndCropJpegFancyUpscaling sets the optional fancy_upscaling attribute to value.
+//
+// value: If true use a slower but nicer upscaling of the
+// chroma planes (yuv420/422 only).
+// If not specified, defaults to true
+func DecodeAndCropJpegFancyUpscaling(value bool) DecodeAndCropJpegAttr {
+ return func(m optionalAttr) {
+ m["fancy_upscaling"] = value
+ }
+}
+
+// DecodeAndCropJpegTryRecoverTruncated sets the optional try_recover_truncated attribute to value.
+//
+// value: If true try to recover an image from truncated input.
+// If not specified, defaults to false
+func DecodeAndCropJpegTryRecoverTruncated(value bool) DecodeAndCropJpegAttr {
+ return func(m optionalAttr) {
+ m["try_recover_truncated"] = value
+ }
+}
+
+// DecodeAndCropJpegAcceptableFraction sets the optional acceptable_fraction attribute to value.
+//
+// value: The minimum required fraction of lines before a truncated
+// input is accepted.
+// If not specified, defaults to 1
+func DecodeAndCropJpegAcceptableFraction(value float32) DecodeAndCropJpegAttr {
+ return func(m optionalAttr) {
+ m["acceptable_fraction"] = value
+ }
+}
+
+// DecodeAndCropJpegDctMethod sets the optional dct_method attribute to value.
+//
+// value: string specifying a hint about the algorithm used for
+// decompression. Defaults to "" which maps to a system-specific
+// default. Currently valid values are ["INTEGER_FAST",
+// "INTEGER_ACCURATE"]. The hint may be ignored (e.g., the internal
+// jpeg library changes to a version that does not have that specific
+// option.)
+// If not specified, defaults to ""
+func DecodeAndCropJpegDctMethod(value string) DecodeAndCropJpegAttr {
+ return func(m optionalAttr) {
+ m["dct_method"] = value
+ }
+}
+
+// Decode and Crop a JPEG-encoded image to a uint8 tensor.
+//
+// The attr `channels` indicates the desired number of color channels for the
+// decoded image.
+//
+// Accepted values are:
+//
+// * 0: Use the number of channels in the JPEG-encoded image.
+// * 1: output a grayscale image.
+// * 3: output an RGB image.
+//
+// If needed, the JPEG-encoded image is transformed to match the requested number
+// of color channels.
+//
+// The attr `ratio` allows downscaling the image by an integer factor during
+// decoding. Allowed values are: 1, 2, 4, and 8. This is much faster than
+// downscaling the image later.
+//
+//
+// It is equivalent to a combination of decode and crop, but much faster by only
+// decoding partial jpeg image.
+//
+// Arguments:
+// contents: 0-D. The JPEG-encoded image.
+// crop_window: 1-D. The crop window: [crop_y, crop_x, crop_height, crop_width].
+//
+// Returns 3-D with shape `[height, width, channels]`..
+func DecodeAndCropJpeg(scope *Scope, contents tf.Output, crop_window tf.Output, optional ...DecodeAndCropJpegAttr) (image tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "DecodeAndCropJpeg",
+ Input: []tf.Input{
+ contents, crop_window,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Adds two `SparseTensor` objects to produce another `SparseTensor`.
+//
+// The input `SparseTensor` objects' indices are assumed ordered in standard
+// lexicographic order. If this is not the case, before this step run
+// `SparseReorder` to restore index ordering.
+//
+// By default, if two values sum to zero at some index, the output `SparseTensor`
+// would still include that particular location in its index, storing a zero in the
+// corresponding value slot. To override this, callers can specify `thresh`,
+// indicating that if the sum has a magnitude strictly smaller than `thresh`, its
+// corresponding value and index would then not be included. In particular,
+// `thresh == 0` (default) means everything is kept and actual thresholding happens
+// only for a positive value.
+//
+// In the following shapes, `nnz` is the count after taking `thresh` into account.
+//
+// Arguments:
+// a_indices: 2-D. The `indices` of the first `SparseTensor`, size `[nnz, ndims]` Matrix.
+// a_values: 1-D. The `values` of the first `SparseTensor`, size `[nnz]` Vector.
+// a_shape: 1-D. The `shape` of the first `SparseTensor`, size `[ndims]` Vector.
+// b_indices: 2-D. The `indices` of the second `SparseTensor`, size `[nnz, ndims]` Matrix.
+// b_values: 1-D. The `values` of the second `SparseTensor`, size `[nnz]` Vector.
+// b_shape: 1-D. The `shape` of the second `SparseTensor`, size `[ndims]` Vector.
+// thresh: 0-D. The magnitude threshold that determines if an output value/index
+// pair takes space.
+func SparseAdd(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b_indices tf.Output, b_values tf.Output, b_shape tf.Output, thresh tf.Output) (sum_indices tf.Output, sum_values tf.Output, sum_shape tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseAdd",
+ Input: []tf.Input{
+ a_indices, a_values, a_shape, b_indices, b_values, b_shape, thresh,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// QuantizedRelu6Attr is an optional argument to QuantizedRelu6.
+type QuantizedRelu6Attr func(optionalAttr)
+
+// QuantizedRelu6OutType sets the optional out_type attribute to value.
+// If not specified, defaults to DT_QUINT8
+func QuantizedRelu6OutType(value tf.DataType) QuantizedRelu6Attr {
+ return func(m optionalAttr) {
+ m["out_type"] = value
+ }
+}
+
+// Computes Quantized Rectified Linear 6: `min(max(features, 0), 6)`
+//
+// Arguments:
+//
+// min_features: The float value that the lowest quantized value represents.
+// max_features: The float value that the highest quantized value represents.
+//
+// Returns Has the same output shape as "features".The float value that the lowest quantized value represents.The float value that the highest quantized value represents.
+func QuantizedRelu6(scope *Scope, features tf.Output, min_features tf.Output, max_features tf.Output, optional ...QuantizedRelu6Attr) (activations tf.Output, min_activations tf.Output, max_activations tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "QuantizedRelu6",
+ Input: []tf.Input{
+ features, min_features, max_features,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// FixedLengthRecordReaderV2Attr is an optional argument to FixedLengthRecordReaderV2.
+type FixedLengthRecordReaderV2Attr func(optionalAttr)
+
+// FixedLengthRecordReaderV2HeaderBytes sets the optional header_bytes attribute to value.
+//
+// value: Number of bytes in the header, defaults to 0.
+// If not specified, defaults to 0
+func FixedLengthRecordReaderV2HeaderBytes(value int64) FixedLengthRecordReaderV2Attr {
+ return func(m optionalAttr) {
+ m["header_bytes"] = value
+ }
+}
+
+// FixedLengthRecordReaderV2FooterBytes sets the optional footer_bytes attribute to value.
+//
+// value: Number of bytes in the footer, defaults to 0.
+// If not specified, defaults to 0
+func FixedLengthRecordReaderV2FooterBytes(value int64) FixedLengthRecordReaderV2Attr {
+ return func(m optionalAttr) {
+ m["footer_bytes"] = value
+ }
+}
+
+// FixedLengthRecordReaderV2HopBytes sets the optional hop_bytes attribute to value.
+//
+// value: Number of bytes to hop before each read. Default of 0 means using
+// record_bytes.
+// If not specified, defaults to 0
+func FixedLengthRecordReaderV2HopBytes(value int64) FixedLengthRecordReaderV2Attr {
+ return func(m optionalAttr) {
+ m["hop_bytes"] = value
+ }
+}
+
+// FixedLengthRecordReaderV2Container sets the optional container attribute to value.
+//
+// value: If non-empty, this reader is placed in the given container.
+// Otherwise, a default container is used.
+// If not specified, defaults to ""
+func FixedLengthRecordReaderV2Container(value string) FixedLengthRecordReaderV2Attr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// FixedLengthRecordReaderV2SharedName sets the optional shared_name attribute to value.
+//
+// value: If non-empty, this reader is named in the given bucket
+// with this shared_name. Otherwise, the node name is used instead.
+// If not specified, defaults to ""
+func FixedLengthRecordReaderV2SharedName(value string) FixedLengthRecordReaderV2Attr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// FixedLengthRecordReaderV2Encoding sets the optional encoding attribute to value.
+//
+// value: The type of encoding for the file. Currently ZLIB and GZIP
+// are supported. Defaults to none.
+// If not specified, defaults to ""
+func FixedLengthRecordReaderV2Encoding(value string) FixedLengthRecordReaderV2Attr {
+ return func(m optionalAttr) {
+ m["encoding"] = value
+ }
+}
+
+// A Reader that outputs fixed-length records from a file.
+//
+// Arguments:
+// record_bytes: Number of bytes in the record.
+//
+// Returns The handle to reference the Reader.
+func FixedLengthRecordReaderV2(scope *Scope, record_bytes int64, optional ...FixedLengthRecordReaderV2Attr) (reader_handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"record_bytes": record_bytes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "FixedLengthRecordReaderV2",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// AudioSummaryAttr is an optional argument to AudioSummary.
type AudioSummaryAttr func(optionalAttr)
@@ -18784,6 +19536,290 @@
return scope.AddOperation(opspec)
}
+// Component-wise divides a SparseTensor by a dense Tensor.
+//
+// *Limitation*: this Op only broadcasts the dense side to the sparse side, but not
+// the other direction.
+//
+// Arguments:
+// sp_indices: 2-D. `N x R` matrix with the indices of non-empty values in a
+// SparseTensor, possibly not in canonical ordering.
+// sp_values: 1-D. `N` non-empty values corresponding to `sp_indices`.
+// sp_shape: 1-D. Shape of the input SparseTensor.
+// dense: `R`-D. The dense Tensor operand.
+//
+// Returns 1-D. The `N` values that are operated on.
+func SparseDenseCwiseDiv(scope *Scope, sp_indices tf.Output, sp_values tf.Output, sp_shape tf.Output, dense tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseDenseCwiseDiv",
+ Input: []tf.Input{
+ sp_indices, sp_values, sp_shape, dense,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// FractionalAvgPoolGradAttr is an optional argument to FractionalAvgPoolGrad.
+type FractionalAvgPoolGradAttr func(optionalAttr)
+
+// FractionalAvgPoolGradOverlapping sets the optional overlapping attribute to value.
+//
+// value: When set to True, it means when pooling, the values at the boundary
+// of adjacent pooling cells are used by both cells. For example:
+//
+// `index 0 1 2 3 4`
+//
+// `value 20 5 16 3 7`
+//
+// If the pooling sequence is [0, 2, 4], then 16, at index 2 will be used twice.
+// The result would be [41/3, 26/3] for fractional avg pooling.
+// If not specified, defaults to false
+func FractionalAvgPoolGradOverlapping(value bool) FractionalAvgPoolGradAttr {
+ return func(m optionalAttr) {
+ m["overlapping"] = value
+ }
+}
+
+// Computes gradient of the FractionalAvgPool function.
+//
+// Unlike FractionalMaxPoolGrad, we don't need to find arg_max for
+// FractionalAvgPoolGrad, we just need to evenly back-propagate each element of
+// out_backprop to those indices that form the same pooling cell. Therefore, we
+// just need to know the shape of original input tensor, instead of the whole
+// tensor.
+//
+// Arguments:
+// orig_input_tensor_shape: Original input tensor shape for `fractional_avg_pool`
+// out_backprop: 4-D with shape `[batch, height, width, channels]`. Gradients
+// w.r.t. the output of `fractional_avg_pool`.
+// row_pooling_sequence: row pooling sequence, form pooling region with
+// col_pooling_sequence.
+// col_pooling_sequence: column pooling sequence, form pooling region with
+// row_pooling sequence.
+//
+// Returns 4-D. Gradients w.r.t. the input of `fractional_avg_pool`.
+func FractionalAvgPoolGrad(scope *Scope, orig_input_tensor_shape tf.Output, out_backprop tf.Output, row_pooling_sequence tf.Output, col_pooling_sequence tf.Output, optional ...FractionalAvgPoolGradAttr) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "FractionalAvgPoolGrad",
+ Input: []tf.Input{
+ orig_input_tensor_shape, out_backprop, row_pooling_sequence, col_pooling_sequence,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Adds up a `SparseTensor` and a dense `Tensor`, producing a dense `Tensor`.
+//
+// This Op does not require `a_indices` be sorted in standard lexicographic order.
+//
+// Arguments:
+// a_indices: 2-D. The `indices` of the `SparseTensor`, with shape `[nnz, ndims]`.
+// a_values: 1-D. The `values` of the `SparseTensor`, with shape `[nnz]`.
+// a_shape: 1-D. The `shape` of the `SparseTensor`, with shape `[ndims]`.
+// b: `ndims`-D Tensor. With shape `a_shape`.
+func SparseTensorDenseAdd(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b tf.Output) (output tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseTensorDenseAdd",
+ Input: []tf.Input{
+ a_indices, a_values, a_shape, b,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// SparseToSparseSetOperationAttr is an optional argument to SparseToSparseSetOperation.
+type SparseToSparseSetOperationAttr func(optionalAttr)
+
+// SparseToSparseSetOperationValidateIndices sets the optional validate_indices attribute to value.
+// If not specified, defaults to true
+func SparseToSparseSetOperationValidateIndices(value bool) SparseToSparseSetOperationAttr {
+ return func(m optionalAttr) {
+ m["validate_indices"] = value
+ }
+}
+
+// Applies set operation along last dimension of 2 `SparseTensor` inputs.
+//
+// See SetOperationOp::SetOperationFromContext for values of `set_operation`.
+//
+// If `validate_indices` is `True`, `SparseToSparseSetOperation` validates the
+// order and range of `set1` and `set2` indices.
+//
+// Input `set1` is a `SparseTensor` represented by `set1_indices`, `set1_values`,
+// and `set1_shape`. For `set1` ranked `n`, 1st `n-1` dimensions must be the same
+// as `set2`. Dimension `n` contains values in a set, duplicates are allowed but
+// ignored.
+//
+// Input `set2` is a `SparseTensor` represented by `set2_indices`, `set2_values`,
+// and `set2_shape`. For `set2` ranked `n`, 1st `n-1` dimensions must be the same
+// as `set1`. Dimension `n` contains values in a set, duplicates are allowed but
+// ignored.
+//
+// If `validate_indices` is `True`, this op validates the order and range of `set1`
+// and `set2` indices.
+//
+// Output `result` is a `SparseTensor` represented by `result_indices`,
+// `result_values`, and `result_shape`. For `set1` and `set2` ranked `n`, this
+// has rank `n` and the same 1st `n-1` dimensions as `set1` and `set2`. The `nth`
+// dimension contains the result of `set_operation` applied to the corresponding
+// `[0...n-1]` dimension of `set`.
+//
+// Arguments:
+// set1_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major
+// order.
+// set1_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major
+// order.
+// set1_shape: 1D `Tensor`, shape of a `SparseTensor`. `set1_shape[0...n-1]` must
+// be the same as `set2_shape[0...n-1]`, `set1_shape[n]` is the
+// max set size across `0...n-1` dimensions.
+// set2_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major
+// order.
+// set2_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major
+// order.
+// set2_shape: 1D `Tensor`, shape of a `SparseTensor`. `set2_shape[0...n-1]` must
+// be the same as `set1_shape[0...n-1]`, `set2_shape[n]` is the
+// max set size across `0...n-1` dimensions.
+//
+//
+// Returns 2D indices of a `SparseTensor`.1D values of a `SparseTensor`.1D `Tensor` shape of a `SparseTensor`. `result_shape[0...n-1]` is
+// the same as the 1st `n-1` dimensions of `set1` and `set2`, `result_shape[n]`
+// is the max result set size across all `0...n-1` dimensions.
+func SparseToSparseSetOperation(scope *Scope, set1_indices tf.Output, set1_values tf.Output, set1_shape tf.Output, set2_indices tf.Output, set2_values tf.Output, set2_shape tf.Output, set_operation string, optional ...SparseToSparseSetOperationAttr) (result_indices tf.Output, result_values tf.Output, result_shape tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"set_operation": set_operation}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "SparseToSparseSetOperation",
+ Input: []tf.Input{
+ set1_indices, set1_values, set1_shape, set2_indices, set2_values, set2_shape,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0), op.Output(1), op.Output(2)
+}
+
+// MutableDenseHashTableV2Attr is an optional argument to MutableDenseHashTableV2.
+type MutableDenseHashTableV2Attr func(optionalAttr)
+
+// MutableDenseHashTableV2Container sets the optional container attribute to value.
+//
+// value: If non-empty, this table is placed in the given container.
+// Otherwise, a default container is used.
+// If not specified, defaults to ""
+func MutableDenseHashTableV2Container(value string) MutableDenseHashTableV2Attr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// MutableDenseHashTableV2SharedName sets the optional shared_name attribute to value.
+//
+// value: If non-empty, this table is shared under the given name across
+// multiple sessions.
+// If not specified, defaults to ""
+func MutableDenseHashTableV2SharedName(value string) MutableDenseHashTableV2Attr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// MutableDenseHashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value.
+// If not specified, defaults to false
+func MutableDenseHashTableV2UseNodeNameSharing(value bool) MutableDenseHashTableV2Attr {
+ return func(m optionalAttr) {
+ m["use_node_name_sharing"] = value
+ }
+}
+
+// MutableDenseHashTableV2ValueShape sets the optional value_shape attribute to value.
+//
+// value: The shape of each value.
+// If not specified, defaults to <>
+func MutableDenseHashTableV2ValueShape(value tf.Shape) MutableDenseHashTableV2Attr {
+ return func(m optionalAttr) {
+ m["value_shape"] = value
+ }
+}
+
+// MutableDenseHashTableV2InitialNumBuckets sets the optional initial_num_buckets attribute to value.
+//
+// value: The initial number of hash table buckets. Must be a power
+// to 2.
+// If not specified, defaults to 131072
+func MutableDenseHashTableV2InitialNumBuckets(value int64) MutableDenseHashTableV2Attr {
+ return func(m optionalAttr) {
+ m["initial_num_buckets"] = value
+ }
+}
+
+// MutableDenseHashTableV2MaxLoadFactor sets the optional max_load_factor attribute to value.
+//
+// value: The maximum ratio between number of entries and number of
+// buckets before growing the table. Must be between 0 and 1.
+// If not specified, defaults to 0.8
+func MutableDenseHashTableV2MaxLoadFactor(value float32) MutableDenseHashTableV2Attr {
+ return func(m optionalAttr) {
+ m["max_load_factor"] = value
+ }
+}
+
+// Creates an empty hash table that uses tensors as the backing store.
+//
+// It uses "open addressing" with quadratic reprobing to resolve
+// collisions.
+//
+// This op creates a mutable hash table, specifying the type of its keys and
+// values. Each value must be a scalar. Data can be inserted into the table using
+// the insert operations. It does not support the initialization operation.
+//
+// Arguments:
+// empty_key: The key used to represent empty key buckets internally. Must not
+// be used in insert or lookup operations.
+//
+// value_dtype: Type of the table values.
+//
+// Returns Handle to a table.
+func MutableDenseHashTableV2(scope *Scope, empty_key tf.Output, deleted_key tf.Output, value_dtype tf.DataType, optional ...MutableDenseHashTableV2Attr) (table_handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"value_dtype": value_dtype}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MutableDenseHashTableV2",
+ Input: []tf.Input{
+ empty_key, deleted_key,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// UpperBoundAttr is an optional argument to UpperBound.
type UpperBoundAttr func(optionalAttr)
@@ -19017,587 +20053,6 @@
return op.Output(0)
}
-// AddManySparseToTensorsMapAttr is an optional argument to AddManySparseToTensorsMap.
-type AddManySparseToTensorsMapAttr func(optionalAttr)
-
-// AddManySparseToTensorsMapContainer sets the optional container attribute to value.
-//
-// value: The container name for the `SparseTensorsMap` created by this op.
-// If not specified, defaults to ""
-func AddManySparseToTensorsMapContainer(value string) AddManySparseToTensorsMapAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// AddManySparseToTensorsMapSharedName sets the optional shared_name attribute to value.
-//
-// value: The shared name for the `SparseTensorsMap` created by this op.
-// If blank, the new Operation's unique name is used.
-// If not specified, defaults to ""
-func AddManySparseToTensorsMapSharedName(value string) AddManySparseToTensorsMapAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Add an `N`-minibatch `SparseTensor` to a `SparseTensorsMap`, return `N` handles.
-//
-// A `SparseTensor` of rank `R` is represented by three tensors: `sparse_indices`,
-// `sparse_values`, and `sparse_shape`, where
-//
-// ```sparse_indices.shape[1] == sparse_shape.shape[0] == R```
-//
-// An `N`-minibatch of `SparseTensor` objects is represented as a `SparseTensor`
-// having a first `sparse_indices` column taking values between `[0, N)`, where
-// the minibatch size `N == sparse_shape[0]`.
-//
-// The input `SparseTensor` must have rank `R` greater than 1, and the first
-// dimension is treated as the minibatch dimension. Elements of the `SparseTensor`
-// must be sorted in increasing order of this first dimension. The stored
-// `SparseTensor` objects pointed to by each row of the output `sparse_handles`
-// will have rank `R-1`.
-//
-// The `SparseTensor` values can then be read out as part of a minibatch by passing
-// the given keys as vector elements to `TakeManySparseFromTensorsMap`. To ensure
-// the correct `SparseTensorsMap` is accessed, ensure that the same
-// `container` and `shared_name` are passed to that Op. If no `shared_name`
-// is provided here, instead use the *name* of the Operation created by calling
-// `AddManySparseToTensorsMap` as the `shared_name` passed to
-// `TakeManySparseFromTensorsMap`. Ensure the Operations are colocated.
-//
-// Arguments:
-// sparse_indices: 2-D. The `indices` of the minibatch `SparseTensor`.
-// `sparse_indices[:, 0]` must be ordered values in `[0, N)`.
-// sparse_values: 1-D. The `values` of the minibatch `SparseTensor`.
-// sparse_shape: 1-D. The `shape` of the minibatch `SparseTensor`.
-// The minibatch size `N == sparse_shape[0]`.
-//
-// Returns 1-D. The handles of the `SparseTensor` now stored in the
-// `SparseTensorsMap`. Shape: `[N]`.
-func AddManySparseToTensorsMap(scope *Scope, sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output, optional ...AddManySparseToTensorsMapAttr) (sparse_handles tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "AddManySparseToTensorsMap",
- Input: []tf.Input{
- sparse_indices, sparse_values, sparse_shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Concatenates tensors along one dimension.
-//
-// Arguments:
-// values: List of `N` Tensors to concatenate. Their ranks and types must match,
-// and their sizes must match in all dimensions except `concat_dim`.
-// axis: 0-D. The dimension along which to concatenate. Must be in the
-// range [-rank(values), rank(values)).
-//
-// Returns A `Tensor` with the concatenation of values stacked along the
-// `concat_dim` dimension. This tensor's shape matches that of `values` except
-// in `concat_dim` where it has the sum of the sizes.
-func ConcatV2(scope *Scope, values []tf.Output, axis tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ConcatV2",
- Input: []tf.Input{
- tf.OutputList(values), axis,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Reads and outputs the entire contents of the input filename.
-func ReadFile(scope *Scope, filename tf.Output) (contents tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ReadFile",
- Input: []tf.Input{
- filename,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Multiplies sparse updates into the variable referenced by `resource`.
-//
-// This operation computes
-//
-// # Scalar indices
-// ref[indices, ...] *= updates[...]
-//
-// # Vector indices (for each i)
-// ref[indices[i], ...] *= updates[i, ...]
-//
-// # High rank indices (for each i, ..., j)
-// ref[indices[i, ..., j], ...] *= updates[i, ..., j, ...]
-//
-// Duplicate entries are handled correctly: if multiple `indices` reference
-// the same location, their contributions multiply.
-//
-// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`.
-//
-// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
-// <img style="width:100%" src='https://www.tensorflow.org/images/ScatterAdd.png' alt>
-// </div>
-//
-// Arguments:
-// resource: Should be from a `Variable` node.
-// indices: A tensor of indices into the first dimension of `ref`.
-// updates: A tensor of updated values to add to `ref`.
-//
-// Returns the created operation.
-func ResourceScatterMul(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ResourceScatterMul",
- Input: []tf.Input{
- resource, indices, updates,
- },
- }
- return scope.AddOperation(opspec)
-}
-
-// Component-wise divides a SparseTensor by a dense Tensor.
-//
-// *Limitation*: this Op only broadcasts the dense side to the sparse side, but not
-// the other direction.
-//
-// Arguments:
-// sp_indices: 2-D. `N x R` matrix with the indices of non-empty values in a
-// SparseTensor, possibly not in canonical ordering.
-// sp_values: 1-D. `N` non-empty values corresponding to `sp_indices`.
-// sp_shape: 1-D. Shape of the input SparseTensor.
-// dense: `R`-D. The dense Tensor operand.
-//
-// Returns 1-D. The `N` values that are operated on.
-func SparseDenseCwiseDiv(scope *Scope, sp_indices tf.Output, sp_values tf.Output, sp_shape tf.Output, dense tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SparseDenseCwiseDiv",
- Input: []tf.Input{
- sp_indices, sp_values, sp_shape, dense,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// FractionalAvgPoolGradAttr is an optional argument to FractionalAvgPoolGrad.
-type FractionalAvgPoolGradAttr func(optionalAttr)
-
-// FractionalAvgPoolGradOverlapping sets the optional overlapping attribute to value.
-//
-// value: When set to True, it means when pooling, the values at the boundary
-// of adjacent pooling cells are used by both cells. For example:
-//
-// `index 0 1 2 3 4`
-//
-// `value 20 5 16 3 7`
-//
-// If the pooling sequence is [0, 2, 4], then 16, at index 2 will be used twice.
-// The result would be [41/3, 26/3] for fractional avg pooling.
-// If not specified, defaults to false
-func FractionalAvgPoolGradOverlapping(value bool) FractionalAvgPoolGradAttr {
- return func(m optionalAttr) {
- m["overlapping"] = value
- }
-}
-
-// Computes gradient of the FractionalAvgPool function.
-//
-// Unlike FractionalMaxPoolGrad, we don't need to find arg_max for
-// FractionalAvgPoolGrad, we just need to evenly back-propagate each element of
-// out_backprop to those indices that form the same pooling cell. Therefore, we
-// just need to know the shape of original input tensor, instead of the whole
-// tensor.
-//
-// Arguments:
-// orig_input_tensor_shape: Original input tensor shape for `fractional_avg_pool`
-// out_backprop: 4-D with shape `[batch, height, width, channels]`. Gradients
-// w.r.t. the output of `fractional_avg_pool`.
-// row_pooling_sequence: row pooling sequence, form pooling region with
-// col_pooling_sequence.
-// col_pooling_sequence: column pooling sequence, form pooling region with
-// row_pooling sequence.
-//
-// Returns 4-D. Gradients w.r.t. the input of `fractional_avg_pool`.
-func FractionalAvgPoolGrad(scope *Scope, orig_input_tensor_shape tf.Output, out_backprop tf.Output, row_pooling_sequence tf.Output, col_pooling_sequence tf.Output, optional ...FractionalAvgPoolGradAttr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "FractionalAvgPoolGrad",
- Input: []tf.Input{
- orig_input_tensor_shape, out_backprop, row_pooling_sequence, col_pooling_sequence,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// SerializeManySparseAttr is an optional argument to SerializeManySparse.
-type SerializeManySparseAttr func(optionalAttr)
-
-// SerializeManySparseOutType sets the optional out_type attribute to value.
-//
-// value: The `dtype` to use for serialization; the supported types are `string`
-// (default) and `variant`.
-// If not specified, defaults to DT_STRING
-func SerializeManySparseOutType(value tf.DataType) SerializeManySparseAttr {
- return func(m optionalAttr) {
- m["out_type"] = value
- }
-}
-
-// Serialize an `N`-minibatch `SparseTensor` into an `[N, 3]` `Tensor` object.
-//
-// The `SparseTensor` must have rank `R` greater than 1, and the first dimension
-// is treated as the minibatch dimension. Elements of the `SparseTensor`
-// must be sorted in increasing order of this first dimension. The serialized
-// `SparseTensor` objects going into each row of `serialized_sparse` will have
-// rank `R-1`.
-//
-// The minibatch size `N` is extracted from `sparse_shape[0]`.
-//
-// Arguments:
-// sparse_indices: 2-D. The `indices` of the minibatch `SparseTensor`.
-// sparse_values: 1-D. The `values` of the minibatch `SparseTensor`.
-// sparse_shape: 1-D. The `shape` of the minibatch `SparseTensor`.
-func SerializeManySparse(scope *Scope, sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output, optional ...SerializeManySparseAttr) (serialized_sparse tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "SerializeManySparse",
- Input: []tf.Input{
- sparse_indices, sparse_values, sparse_shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// Computes inverse hyperbolic cosine of x element-wise.
-func Acosh(scope *Scope, x tf.Output) (y tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "Acosh",
- Input: []tf.Input{
- x,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// TensorArrayV2Attr is an optional argument to TensorArrayV2.
-type TensorArrayV2Attr func(optionalAttr)
-
-// TensorArrayV2ElementShape sets the optional element_shape attribute to value.
-// If not specified, defaults to <unknown_rank:true >
-func TensorArrayV2ElementShape(value tf.Shape) TensorArrayV2Attr {
- return func(m optionalAttr) {
- m["element_shape"] = value
- }
-}
-
-// TensorArrayV2DynamicSize sets the optional dynamic_size attribute to value.
-// If not specified, defaults to false
-func TensorArrayV2DynamicSize(value bool) TensorArrayV2Attr {
- return func(m optionalAttr) {
- m["dynamic_size"] = value
- }
-}
-
-// TensorArrayV2ClearAfterRead sets the optional clear_after_read attribute to value.
-// If not specified, defaults to true
-func TensorArrayV2ClearAfterRead(value bool) TensorArrayV2Attr {
- return func(m optionalAttr) {
- m["clear_after_read"] = value
- }
-}
-
-// TensorArrayV2TensorArrayName sets the optional tensor_array_name attribute to value.
-// If not specified, defaults to ""
-func TensorArrayV2TensorArrayName(value string) TensorArrayV2Attr {
- return func(m optionalAttr) {
- m["tensor_array_name"] = value
- }
-}
-
-// Deprecated. Use TensorArrayV3
-//
-// DEPRECATED at GraphDef version 26: Use TensorArrayV3
-func TensorArrayV2(scope *Scope, size tf.Output, dtype tf.DataType, optional ...TensorArrayV2Attr) (handle tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtype": dtype}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "TensorArrayV2",
- Input: []tf.Input{
- size,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// ThreadUnsafeUnigramCandidateSamplerAttr is an optional argument to ThreadUnsafeUnigramCandidateSampler.
-type ThreadUnsafeUnigramCandidateSamplerAttr func(optionalAttr)
-
-// ThreadUnsafeUnigramCandidateSamplerSeed sets the optional seed attribute to value.
-//
-// value: If either seed or seed2 are set to be non-zero, the random number
-// generator is seeded by the given seed. Otherwise, it is seeded by a
-// random seed.
-// If not specified, defaults to 0
-func ThreadUnsafeUnigramCandidateSamplerSeed(value int64) ThreadUnsafeUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed"] = value
- }
-}
-
-// ThreadUnsafeUnigramCandidateSamplerSeed2 sets the optional seed2 attribute to value.
-//
-// value: An second seed to avoid seed collision.
-// If not specified, defaults to 0
-func ThreadUnsafeUnigramCandidateSamplerSeed2(value int64) ThreadUnsafeUnigramCandidateSamplerAttr {
- return func(m optionalAttr) {
- m["seed2"] = value
- }
-}
-
-// Generates labels for candidate sampling with a learned unigram distribution.
-//
-// See explanations of candidate sampling and the data formats at
-// go/candidate-sampling.
-//
-// For each batch, this op picks a single set of sampled candidate labels.
-//
-// The advantages of sampling candidates per-batch are simplicity and the
-// possibility of efficient dense matrix multiplication. The disadvantage is that
-// the sampled candidates must be chosen independently of the context and of the
-// true labels.
-//
-// Arguments:
-// true_classes: A batch_size * num_true matrix, in which each row contains the
-// IDs of the num_true target_classes in the corresponding original label.
-// num_true: Number of true labels per context.
-// num_sampled: Number of candidates to randomly sample.
-// unique: If unique is true, we sample with rejection, so that all sampled
-// candidates in a batch are unique. This requires some approximation to
-// estimate the post-rejection sampling probabilities.
-// range_max: The sampler will sample integers from the interval [0, range_max).
-//
-// Returns A vector of length num_sampled, in which each element is
-// the ID of a sampled candidate.A batch_size * num_true matrix, representing
-// the number of times each candidate is expected to occur in a batch
-// of sampled candidates. If unique=true, then this is a probability.A vector of length num_sampled, for each sampled
-// candidate representing the number of times the candidate is expected
-// to occur in a batch of sampled candidates. If unique=true, then this is a
-// probability.
-func ThreadUnsafeUnigramCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, range_max int64, optional ...ThreadUnsafeUnigramCandidateSamplerAttr) (sampled_candidates tf.Output, true_expected_count tf.Output, sampled_expected_count tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"num_true": num_true, "num_sampled": num_sampled, "unique": unique, "range_max": range_max}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "ThreadUnsafeUnigramCandidateSampler",
- Input: []tf.Input{
- true_classes,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
-// MaxPoolV2Attr is an optional argument to MaxPoolV2.
-type MaxPoolV2Attr func(optionalAttr)
-
-// MaxPoolV2DataFormat sets the optional data_format attribute to value.
-//
-// value: Specify the data format of the input and output data. With the
-// default format "NHWC", the data is stored in the order of:
-// [batch, in_height, in_width, in_channels].
-// Alternatively, the format could be "NCHW", the data storage order of:
-// [batch, in_channels, in_height, in_width].
-// If not specified, defaults to "NHWC"
-func MaxPoolV2DataFormat(value string) MaxPoolV2Attr {
- return func(m optionalAttr) {
- m["data_format"] = value
- }
-}
-
-// Performs max pooling on the input.
-//
-// Arguments:
-// input: 4-D input to pool over.
-// ksize: The size of the window for each dimension of the input tensor.
-// strides: The stride of the sliding window for each dimension of the
-// input tensor.
-// padding: The type of padding algorithm to use.
-//
-// Returns The max pooled output tensor.
-func MaxPoolV2(scope *Scope, input tf.Output, ksize tf.Output, strides tf.Output, padding string, optional ...MaxPoolV2Attr) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"padding": padding}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MaxPoolV2",
- Input: []tf.Input{
- input, ksize, strides,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// MutableDenseHashTableV2Attr is an optional argument to MutableDenseHashTableV2.
-type MutableDenseHashTableV2Attr func(optionalAttr)
-
-// MutableDenseHashTableV2Container sets the optional container attribute to value.
-//
-// value: If non-empty, this table is placed in the given container.
-// Otherwise, a default container is used.
-// If not specified, defaults to ""
-func MutableDenseHashTableV2Container(value string) MutableDenseHashTableV2Attr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// MutableDenseHashTableV2SharedName sets the optional shared_name attribute to value.
-//
-// value: If non-empty, this table is shared under the given name across
-// multiple sessions.
-// If not specified, defaults to ""
-func MutableDenseHashTableV2SharedName(value string) MutableDenseHashTableV2Attr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// MutableDenseHashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value.
-// If not specified, defaults to false
-func MutableDenseHashTableV2UseNodeNameSharing(value bool) MutableDenseHashTableV2Attr {
- return func(m optionalAttr) {
- m["use_node_name_sharing"] = value
- }
-}
-
-// MutableDenseHashTableV2ValueShape sets the optional value_shape attribute to value.
-//
-// value: The shape of each value.
-// If not specified, defaults to <>
-func MutableDenseHashTableV2ValueShape(value tf.Shape) MutableDenseHashTableV2Attr {
- return func(m optionalAttr) {
- m["value_shape"] = value
- }
-}
-
-// MutableDenseHashTableV2InitialNumBuckets sets the optional initial_num_buckets attribute to value.
-//
-// value: The initial number of hash table buckets. Must be a power
-// to 2.
-// If not specified, defaults to 131072
-func MutableDenseHashTableV2InitialNumBuckets(value int64) MutableDenseHashTableV2Attr {
- return func(m optionalAttr) {
- m["initial_num_buckets"] = value
- }
-}
-
-// MutableDenseHashTableV2MaxLoadFactor sets the optional max_load_factor attribute to value.
-//
-// value: The maximum ratio between number of entries and number of
-// buckets before growing the table. Must be between 0 and 1.
-// If not specified, defaults to 0.8
-func MutableDenseHashTableV2MaxLoadFactor(value float32) MutableDenseHashTableV2Attr {
- return func(m optionalAttr) {
- m["max_load_factor"] = value
- }
-}
-
-// Creates an empty hash table that uses tensors as the backing store.
-//
-// It uses "open addressing" with quadratic reprobing to resolve
-// collisions.
-//
-// This op creates a mutable hash table, specifying the type of its keys and
-// values. Each value must be a scalar. Data can be inserted into the table using
-// the insert operations. It does not support the initialization operation.
-//
-// Arguments:
-// empty_key: The key used to represent empty key buckets internally. Must not
-// be used in insert or lookup operations.
-//
-// value_dtype: Type of the table values.
-//
-// Returns Handle to a table.
-func MutableDenseHashTableV2(scope *Scope, empty_key tf.Output, deleted_key tf.Output, value_dtype tf.DataType, optional ...MutableDenseHashTableV2Attr) (table_handle tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"value_dtype": value_dtype}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MutableDenseHashTableV2",
- Input: []tf.Input{
- empty_key, deleted_key,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Computes the gradient for the inverse of `x` wrt its input.
//
// Specifically, `grad = -dy * y*y`, where `y = 1/x`, and `dy`
@@ -19993,58 +20448,61 @@
return op.Output(0)
}
-// StageSizeAttr is an optional argument to StageSize.
-type StageSizeAttr func(optionalAttr)
-
-// StageSizeCapacity sets the optional capacity attribute to value.
-// If not specified, defaults to 0
+// Inverse 3D real-valued fast Fourier transform.
//
-// REQUIRES: value >= 0
-func StageSizeCapacity(value int64) StageSizeAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// StageSizeMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
+// Computes the inverse 3-dimensional discrete Fourier transform of a real-valued
+// signal over the inner-most 3 dimensions of `input`.
//
-// REQUIRES: value >= 0
-func StageSizeMemoryLimit(value int64) StageSizeAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// StageSizeContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func StageSizeContainer(value string) StageSizeAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// StageSizeSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func StageSizeSharedName(value string) StageSizeAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op returns the number of elements in the underlying container.
-func StageSize(scope *Scope, dtypes []tf.DataType, optional ...StageSizeAttr) (size tf.Output) {
+// The inner-most 3 dimensions of `input` are assumed to be the result of `RFFT3D`:
+// The inner-most dimension contains the `fft_length / 2 + 1` unique components of
+// the DFT of a real-valued signal. If `fft_length` is not provided, it is computed
+// from the size of the inner-most 3 dimensions of `input`. If the FFT length used
+// to compute `input` is odd, it should be provided since it cannot be inferred
+// properly.
+//
+// Along each axis `IRFFT3D` is computed on, if `fft_length` (or
+// `fft_length / 2 + 1` for the inner-most dimension) is smaller than the
+// corresponding dimension of `input`, the dimension is cropped. If it is larger,
+// the dimension is padded with zeros.
+//
+// Arguments:
+// input: A complex64 tensor.
+// fft_length: An int32 tensor of shape [3]. The FFT length for each dimension.
+//
+// Returns A float32 tensor of the same rank as `input`. The inner-most 3
+// dimensions of `input` are replaced with the `fft_length` samples of their
+// inverse 3D real Fourier transform.
+//
+// @compatibility(numpy)
+// Equivalent to np.irfftn with 3 dimensions.
+// @end_compatibility
+func IRFFT3D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) {
if scope.Err() != nil {
return
}
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
+ opspec := tf.OpSpec{
+ Type: "IRFFT3D",
+ Input: []tf.Input{
+ input, fft_length,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// Returns the truth value of (x != y) element-wise.
+//
+// *NOTE*: `NotEqual` supports broadcasting. More about broadcasting
+// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+func NotEqual(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
+ if scope.Err() != nil {
+ return
}
opspec := tf.OpSpec{
- Type: "StageSize",
-
- Attrs: attrs,
+ Type: "NotEqual",
+ Input: []tf.Input{
+ x, y,
+ },
}
op := scope.AddOperation(opspec)
return op.Output(0)
@@ -20212,24 +20670,6 @@
return op.Output(0)
}
-// Returns the truth value of (x >= y) element-wise.
-//
-// *NOTE*: `GreaterEqual` supports broadcasting. More about broadcasting
-// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
-func GreaterEqual(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "GreaterEqual",
- Input: []tf.Input{
- x, y,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// Conv3DAttr is an optional argument to Conv3D.
type Conv3DAttr func(optionalAttr)
@@ -20295,39 +20735,6 @@
return op.Output(0)
}
-// Adds up a SparseTensor and a dense Tensor, using these special rules:
-//
-// (1) Broadcasts the dense side to have the same shape as the sparse side, if
-// eligible;
-// (2) Then, only the dense values pointed to by the indices of the SparseTensor
-// participate in the cwise addition.
-//
-// By these rules, the result is a logical SparseTensor with exactly the same
-// indices and shape, but possibly with different non-zero values. The output of
-// this Op is the resultant non-zero values.
-//
-// Arguments:
-// sp_indices: 2-D. `N x R` matrix with the indices of non-empty values in a
-// SparseTensor, possibly not in canonical ordering.
-// sp_values: 1-D. `N` non-empty values corresponding to `sp_indices`.
-// sp_shape: 1-D. Shape of the input SparseTensor.
-// dense: `R`-D. The dense Tensor operand.
-//
-// Returns 1-D. The `N` values that are operated on.
-func SparseDenseCwiseAdd(scope *Scope, sp_indices tf.Output, sp_values tf.Output, sp_shape tf.Output, dense tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "SparseDenseCwiseAdd",
- Input: []tf.Input{
- sp_indices, sp_values, sp_shape, dense,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// QuantizeV2Attr is an optional argument to QuantizeV2.
type QuantizeV2Attr func(optionalAttr)
@@ -21394,6 +21801,142 @@
return op.Output(0), op.Output(1)
}
+// MapUnstageNoKeyAttr is an optional argument to MapUnstageNoKey.
+type MapUnstageNoKeyAttr func(optionalAttr)
+
+// MapUnstageNoKeyCapacity sets the optional capacity attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapUnstageNoKeyCapacity(value int64) MapUnstageNoKeyAttr {
+ return func(m optionalAttr) {
+ m["capacity"] = value
+ }
+}
+
+// MapUnstageNoKeyMemoryLimit sets the optional memory_limit attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func MapUnstageNoKeyMemoryLimit(value int64) MapUnstageNoKeyAttr {
+ return func(m optionalAttr) {
+ m["memory_limit"] = value
+ }
+}
+
+// MapUnstageNoKeyContainer sets the optional container attribute to value.
+// If not specified, defaults to ""
+func MapUnstageNoKeyContainer(value string) MapUnstageNoKeyAttr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// MapUnstageNoKeySharedName sets the optional shared_name attribute to value.
+// If not specified, defaults to ""
+func MapUnstageNoKeySharedName(value string) MapUnstageNoKeyAttr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// Op removes and returns a random (key, value)
+//
+// from the underlying container. If the underlying container
+// does not contain elements, the op will block until it does.
+func MapUnstageNoKey(scope *Scope, indices tf.Output, dtypes []tf.DataType, optional ...MapUnstageNoKeyAttr) (key tf.Output, values []tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"dtypes": dtypes}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "MapUnstageNoKey",
+ Input: []tf.Input{
+ indices,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ if scope.Err() != nil {
+ return
+ }
+ var idx int
+ var err error
+ key = op.Output(idx)
+ if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
+ scope.UpdateErr("MapUnstageNoKey", err)
+ return
+ }
+ return key, values
+}
+
+// HashTableV2Attr is an optional argument to HashTableV2.
+type HashTableV2Attr func(optionalAttr)
+
+// HashTableV2Container sets the optional container attribute to value.
+//
+// value: If non-empty, this table is placed in the given container.
+// Otherwise, a default container is used.
+// If not specified, defaults to ""
+func HashTableV2Container(value string) HashTableV2Attr {
+ return func(m optionalAttr) {
+ m["container"] = value
+ }
+}
+
+// HashTableV2SharedName sets the optional shared_name attribute to value.
+//
+// value: If non-empty, this table is shared under the given name across
+// multiple sessions.
+// If not specified, defaults to ""
+func HashTableV2SharedName(value string) HashTableV2Attr {
+ return func(m optionalAttr) {
+ m["shared_name"] = value
+ }
+}
+
+// HashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value.
+//
+// value: If true and shared_name is empty, the table is shared
+// using the node name.
+// If not specified, defaults to false
+func HashTableV2UseNodeNameSharing(value bool) HashTableV2Attr {
+ return func(m optionalAttr) {
+ m["use_node_name_sharing"] = value
+ }
+}
+
+// Creates a non-initialized hash table.
+//
+// This op creates a hash table, specifying the type of its keys and values.
+// Before using the table you will have to initialize it. After initialization the
+// table will be immutable.
+//
+// Arguments:
+// key_dtype: Type of the table keys.
+// value_dtype: Type of the table values.
+//
+// Returns Handle to a table.
+func HashTableV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...HashTableV2Attr) (table_handle tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "HashTableV2",
+
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// TakeManySparseFromTensorsMapAttr is an optional argument to TakeManySparseFromTensorsMap.
type TakeManySparseFromTensorsMapAttr func(optionalAttr)
@@ -21773,93 +22316,6 @@
return scope.AddOperation(opspec)
}
-// L2 Loss.
-//
-// Computes half the L2 norm of a tensor without the `sqrt`:
-//
-// output = sum(t ** 2) / 2
-//
-// Arguments:
-// t: Typically 2-D, but may have any dimensions.
-//
-// Returns 0-D.
-func L2Loss(scope *Scope, t tf.Output) (output tf.Output) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "L2Loss",
- Input: []tf.Input{
- t,
- },
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
-// DenseToSparseSetOperationAttr is an optional argument to DenseToSparseSetOperation.
-type DenseToSparseSetOperationAttr func(optionalAttr)
-
-// DenseToSparseSetOperationValidateIndices sets the optional validate_indices attribute to value.
-// If not specified, defaults to true
-func DenseToSparseSetOperationValidateIndices(value bool) DenseToSparseSetOperationAttr {
- return func(m optionalAttr) {
- m["validate_indices"] = value
- }
-}
-
-// Applies set operation along last dimension of `Tensor` and `SparseTensor`.
-//
-// See SetOperationOp::SetOperationFromContext for values of `set_operation`.
-//
-// Input `set2` is a `SparseTensor` represented by `set2_indices`, `set2_values`,
-// and `set2_shape`. For `set2` ranked `n`, 1st `n-1` dimensions must be the same
-// as `set1`. Dimension `n` contains values in a set, duplicates are allowed but
-// ignored.
-//
-// If `validate_indices` is `True`, this op validates the order and range of `set2`
-// indices.
-//
-// Output `result` is a `SparseTensor` represented by `result_indices`,
-// `result_values`, and `result_shape`. For `set1` and `set2` ranked `n`, this
-// has rank `n` and the same 1st `n-1` dimensions as `set1` and `set2`. The `nth`
-// dimension contains the result of `set_operation` applied to the corresponding
-// `[0...n-1]` dimension of `set`.
-//
-// Arguments:
-// set1: `Tensor` with rank `n`. 1st `n-1` dimensions must be the same as `set2`.
-// Dimension `n` contains values in a set, duplicates are allowed but ignored.
-// set2_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major
-// order.
-// set2_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major
-// order.
-// set2_shape: 1D `Tensor`, shape of a `SparseTensor`. `set2_shape[0...n-1]` must
-// be the same as the 1st `n-1` dimensions of `set1`, `result_shape[n]` is the
-// max set size across `n-1` dimensions.
-//
-//
-// Returns 2D indices of a `SparseTensor`.1D values of a `SparseTensor`.1D `Tensor` shape of a `SparseTensor`. `result_shape[0...n-1]` is
-// the same as the 1st `n-1` dimensions of `set1` and `set2`, `result_shape[n]`
-// is the max result set size across all `0...n-1` dimensions.
-func DenseToSparseSetOperation(scope *Scope, set1 tf.Output, set2_indices tf.Output, set2_values tf.Output, set2_shape tf.Output, set_operation string, optional ...DenseToSparseSetOperationAttr) (result_indices tf.Output, result_values tf.Output, result_shape tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"set_operation": set_operation}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "DenseToSparseSetOperation",
- Input: []tf.Input{
- set1, set2_indices, set2_values, set2_shape,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0), op.Output(1), op.Output(2)
-}
-
// Subtracts a value from the current value of a variable.
//
// Any ReadVariableOp with a control dependency on this op is guaranteed to
@@ -22388,47 +22844,6 @@
return output
}
-// Adds sparse updates to the variable referenced by `resource`.
-//
-// This operation computes
-//
-// # Scalar indices
-// ref[indices, ...] += updates[...]
-//
-// # Vector indices (for each i)
-// ref[indices[i], ...] += updates[i, ...]
-//
-// # High rank indices (for each i, ..., j)
-// ref[indices[i, ..., j], ...] += updates[i, ..., j, ...]
-//
-// Duplicate entries are handled correctly: if multiple `indices` reference
-// the same location, their contributions add.
-//
-// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`.
-//
-// <div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
-// <img style="width:100%" src='https://www.tensorflow.org/images/ScatterAdd.png' alt>
-// </div>
-//
-// Arguments:
-// resource: Should be from a `Variable` node.
-// indices: A tensor of indices into the first dimension of `ref`.
-// updates: A tensor of updated values to add to `ref`.
-//
-// Returns the created operation.
-func ResourceScatterAdd(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- opspec := tf.OpSpec{
- Type: "ResourceScatterAdd",
- Input: []tf.Input{
- resource, indices, updates,
- },
- }
- return scope.AddOperation(opspec)
-}
-
// Says whether the targets are in the top `K` predictions.
//
// This outputs a `batch_size` bool array, an entry `out[i]` is `true` if the
@@ -22593,6 +23008,55 @@
return op.Output(0), op.Output(1), op.Output(2)
}
+// Returns the truth value of (x >= y) element-wise.
+//
+// *NOTE*: `GreaterEqual` supports broadcasting. More about broadcasting
+// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+func GreaterEqual(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "GreaterEqual",
+ Input: []tf.Input{
+ x, y,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
+// ApproximateEqualAttr is an optional argument to ApproximateEqual.
+type ApproximateEqualAttr func(optionalAttr)
+
+// ApproximateEqualTolerance sets the optional tolerance attribute to value.
+// If not specified, defaults to 1e-05
+func ApproximateEqualTolerance(value float32) ApproximateEqualAttr {
+ return func(m optionalAttr) {
+ m["tolerance"] = value
+ }
+}
+
+// Returns the truth value of abs(x-y) < tolerance element-wise.
+func ApproximateEqual(scope *Scope, x tf.Output, y tf.Output, optional ...ApproximateEqualAttr) (z tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ attrs := map[string]interface{}{}
+ for _, a := range optional {
+ a(attrs)
+ }
+ opspec := tf.OpSpec{
+ Type: "ApproximateEqual",
+ Input: []tf.Input{
+ x, y,
+ },
+ Attrs: attrs,
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Returns the truth value of x OR y element-wise.
//
// *NOTE*: `LogicalOr` supports broadcasting. More about broadcasting
@@ -25426,77 +25890,6 @@
return op.Output(0)
}
-// MapPeekAttr is an optional argument to MapPeek.
-type MapPeekAttr func(optionalAttr)
-
-// MapPeekCapacity sets the optional capacity attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapPeekCapacity(value int64) MapPeekAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// MapPeekMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapPeekMemoryLimit(value int64) MapPeekAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// MapPeekContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func MapPeekContainer(value string) MapPeekAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// MapPeekSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func MapPeekSharedName(value string) MapPeekAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op peeks at the values at the specified key. If the
-//
-// underlying container does not contain this key
-// this op will block until it does.
-func MapPeek(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.DataType, optional ...MapPeekAttr) (values []tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MapPeek",
- Input: []tf.Input{
- key, indices,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- if scope.Err() != nil {
- return
- }
- var idx int
- var err error
- if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
- scope.UpdateErr("MapPeek", err)
- return
- }
- return values
-}
-
// Looks up keys in a table, outputs the corresponding values.
//
// The tensor `keys` must of the same type as the keys of the table.
@@ -25707,142 +26100,6 @@
return scope.AddOperation(opspec)
}
-// MapUnstageNoKeyAttr is an optional argument to MapUnstageNoKey.
-type MapUnstageNoKeyAttr func(optionalAttr)
-
-// MapUnstageNoKeyCapacity sets the optional capacity attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapUnstageNoKeyCapacity(value int64) MapUnstageNoKeyAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// MapUnstageNoKeyMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapUnstageNoKeyMemoryLimit(value int64) MapUnstageNoKeyAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// MapUnstageNoKeyContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func MapUnstageNoKeyContainer(value string) MapUnstageNoKeyAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// MapUnstageNoKeySharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func MapUnstageNoKeySharedName(value string) MapUnstageNoKeyAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op removes and returns a random (key, value)
-//
-// from the underlying container. If the underlying container
-// does not contain elements, the op will block until it does.
-func MapUnstageNoKey(scope *Scope, indices tf.Output, dtypes []tf.DataType, optional ...MapUnstageNoKeyAttr) (key tf.Output, values []tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MapUnstageNoKey",
- Input: []tf.Input{
- indices,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- if scope.Err() != nil {
- return
- }
- var idx int
- var err error
- key = op.Output(idx)
- if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
- scope.UpdateErr("MapUnstageNoKey", err)
- return
- }
- return key, values
-}
-
-// HashTableV2Attr is an optional argument to HashTableV2.
-type HashTableV2Attr func(optionalAttr)
-
-// HashTableV2Container sets the optional container attribute to value.
-//
-// value: If non-empty, this table is placed in the given container.
-// Otherwise, a default container is used.
-// If not specified, defaults to ""
-func HashTableV2Container(value string) HashTableV2Attr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// HashTableV2SharedName sets the optional shared_name attribute to value.
-//
-// value: If non-empty, this table is shared under the given name across
-// multiple sessions.
-// If not specified, defaults to ""
-func HashTableV2SharedName(value string) HashTableV2Attr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// HashTableV2UseNodeNameSharing sets the optional use_node_name_sharing attribute to value.
-//
-// value: If true and shared_name is empty, the table is shared
-// using the node name.
-// If not specified, defaults to false
-func HashTableV2UseNodeNameSharing(value bool) HashTableV2Attr {
- return func(m optionalAttr) {
- m["use_node_name_sharing"] = value
- }
-}
-
-// Creates a non-initialized hash table.
-//
-// This op creates a hash table, specifying the type of its keys and values.
-// Before using the table you will have to initialize it. After initialization the
-// table will be immutable.
-//
-// Arguments:
-// key_dtype: Type of the table keys.
-// value_dtype: Type of the table values.
-//
-// Returns Handle to a table.
-func HashTableV2(scope *Scope, key_dtype tf.DataType, value_dtype tf.DataType, optional ...HashTableV2Attr) (table_handle tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"key_dtype": key_dtype, "value_dtype": value_dtype}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "HashTableV2",
-
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
-
// MultiDeviceIteratorFromStringHandleAttr is an optional argument to MultiDeviceIteratorFromStringHandle.
type MultiDeviceIteratorFromStringHandleAttr func(optionalAttr)
@@ -31347,6 +31604,24 @@
return op.Output(0)
}
+// Returns the truth value of (x > y) element-wise.
+//
+// *NOTE*: `Greater` supports broadcasting. More about broadcasting
+// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+func Greater(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
+ if scope.Err() != nil {
+ return
+ }
+ opspec := tf.OpSpec{
+ Type: "Greater",
+ Input: []tf.Input{
+ x, y,
+ },
+ }
+ op := scope.AddOperation(opspec)
+ return op.Output(0)
+}
+
// Performs a padding as a preprocess during a convolution.
//
// Similar to FusedResizeAndPadConv2d, this op allows for an optimized
@@ -33395,278 +33670,3 @@
}
return scope.AddOperation(opspec)
}
-
-// StagePeekAttr is an optional argument to StagePeek.
-type StagePeekAttr func(optionalAttr)
-
-// StagePeekCapacity sets the optional capacity attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func StagePeekCapacity(value int64) StagePeekAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// StagePeekMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func StagePeekMemoryLimit(value int64) StagePeekAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// StagePeekContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func StagePeekContainer(value string) StagePeekAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// StagePeekSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func StagePeekSharedName(value string) StagePeekAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op peeks at the values at the specified index. If the
-//
-// underlying container does not contain sufficient elements
-// this op will block until it does. This Op is optimized for
-// performance.
-func StagePeek(scope *Scope, index tf.Output, dtypes []tf.DataType, optional ...StagePeekAttr) (values []tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "StagePeek",
- Input: []tf.Input{
- index,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- if scope.Err() != nil {
- return
- }
- var idx int
- var err error
- if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
- scope.UpdateErr("StagePeek", err)
- return
- }
- return values
-}
-
-// MapStageAttr is an optional argument to MapStage.
-type MapStageAttr func(optionalAttr)
-
-// MapStageCapacity sets the optional capacity attribute to value.
-//
-// value: Maximum number of elements in the Staging Area. If > 0, inserts
-// on the container will block when the capacity is reached.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapStageCapacity(value int64) MapStageAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// MapStageMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapStageMemoryLimit(value int64) MapStageAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// MapStageContainer sets the optional container attribute to value.
-//
-// value: If non-empty, this queue is placed in the given container. Otherwise,
-// a default container is used.
-// If not specified, defaults to ""
-func MapStageContainer(value string) MapStageAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// MapStageSharedName sets the optional shared_name attribute to value.
-//
-// value: It is necessary to match this name to the matching Unstage Op.
-// If not specified, defaults to ""
-func MapStageSharedName(value string) MapStageAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Stage (key, values) in the underlying container which behaves like a hashtable.
-//
-// Arguments:
-// key: int64
-//
-// values: a list of tensors
-// dtypes A list of data types that inserted values should adhere to.
-//
-//
-// Returns the created operation.
-func MapStage(scope *Scope, key tf.Output, indices tf.Output, values []tf.Output, dtypes []tf.DataType, optional ...MapStageAttr) (o *tf.Operation) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MapStage",
- Input: []tf.Input{
- key, indices, tf.OutputList(values),
- },
- Attrs: attrs,
- }
- return scope.AddOperation(opspec)
-}
-
-// MapUnstageAttr is an optional argument to MapUnstage.
-type MapUnstageAttr func(optionalAttr)
-
-// MapUnstageCapacity sets the optional capacity attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapUnstageCapacity(value int64) MapUnstageAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// MapUnstageMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapUnstageMemoryLimit(value int64) MapUnstageAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// MapUnstageContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func MapUnstageContainer(value string) MapUnstageAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// MapUnstageSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func MapUnstageSharedName(value string) MapUnstageAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op removes and returns the values associated with the key
-//
-// from the underlying container. If the underlying container
-// does not contain this key, the op will block until it does.
-func MapUnstage(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.DataType, optional ...MapUnstageAttr) (values []tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MapUnstage",
- Input: []tf.Input{
- key, indices,
- },
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- if scope.Err() != nil {
- return
- }
- var idx int
- var err error
- if values, idx, err = makeOutputList(op, idx, "values"); err != nil {
- scope.UpdateErr("MapUnstage", err)
- return
- }
- return values
-}
-
-// MapSizeAttr is an optional argument to MapSize.
-type MapSizeAttr func(optionalAttr)
-
-// MapSizeCapacity sets the optional capacity attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapSizeCapacity(value int64) MapSizeAttr {
- return func(m optionalAttr) {
- m["capacity"] = value
- }
-}
-
-// MapSizeMemoryLimit sets the optional memory_limit attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func MapSizeMemoryLimit(value int64) MapSizeAttr {
- return func(m optionalAttr) {
- m["memory_limit"] = value
- }
-}
-
-// MapSizeContainer sets the optional container attribute to value.
-// If not specified, defaults to ""
-func MapSizeContainer(value string) MapSizeAttr {
- return func(m optionalAttr) {
- m["container"] = value
- }
-}
-
-// MapSizeSharedName sets the optional shared_name attribute to value.
-// If not specified, defaults to ""
-func MapSizeSharedName(value string) MapSizeAttr {
- return func(m optionalAttr) {
- m["shared_name"] = value
- }
-}
-
-// Op returns the number of elements in the underlying container.
-func MapSize(scope *Scope, dtypes []tf.DataType, optional ...MapSizeAttr) (size tf.Output) {
- if scope.Err() != nil {
- return
- }
- attrs := map[string]interface{}{"dtypes": dtypes}
- for _, a := range optional {
- a(attrs)
- }
- opspec := tf.OpSpec{
- Type: "MapSize",
-
- Attrs: attrs,
- }
- op := scope.AddOperation(opspec)
- return op.Output(0)
-}
diff --git a/tensorflow/go/session.go b/tensorflow/go/session.go
index db6ae4f..48909ff 100644
--- a/tensorflow/go/session.go
+++ b/tensorflow/go/session.go
@@ -71,37 +71,39 @@
MemoryLimitBytes int64
}
-// Return list of devices associated with a Session
-func (s *Session) ListDevices() ([]Device, error) {
- var devices []Device
-
- status := newStatus()
- devices_list := C.TF_SessionListDevices(s.c, status.c)
- if err := status.Err(); err != nil {
- return nil, fmt.Errorf("SessionListDevices() failed: %v", err)
+// String describes d and implements fmt.Stringer.
+func (d Device) String() string {
+ memStr := "no memory limit"
+ if d.MemoryLimitBytes >= 0 {
+ memStr = fmt.Sprintf("memory limit %d bytes", d.MemoryLimitBytes)
}
- defer C.TF_DeleteDeviceList(devices_list)
+ return fmt.Sprintf("(Device: name \"%s\", type %s, %s)", d.Name, d.Type, memStr)
+}
- for i := 0; i < int(C.TF_DeviceListCount(devices_list)); i++ {
- device_name := C.TF_DeviceListName(devices_list, C.int(i), status.c)
+func deviceSliceFromDeviceList(list *C.TF_DeviceList) ([]Device, error) {
+ var devices []Device
+ status := newStatus()
+
+ for i := 0; i < int(C.TF_DeviceListCount(list)); i++ {
+ name := C.TF_DeviceListName(list, C.int(i), status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("DeviceListName(index=%d) failed: %v", i, err)
}
- device_type := C.TF_DeviceListType(devices_list, C.int(i), status.c)
+ deviceType := C.TF_DeviceListType(list, C.int(i), status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("DeviceListType(index=%d) failed: %v", i, err)
}
- memory_limit_bytes := C.TF_DeviceListMemoryBytes(devices_list, C.int(i), status.c)
+ memoryLimitBytes := C.TF_DeviceListMemoryBytes(list, C.int(i), status.c)
if err := status.Err(); err != nil {
return nil, fmt.Errorf("DeviceListMemoryBytes(index=%d) failed: %v", i, err)
}
device := Device{
- Name: C.GoString(device_name),
- Type: C.GoString(device_type),
- MemoryLimitBytes: int64(memory_limit_bytes),
+ Name: C.GoString(name),
+ Type: C.GoString(deviceType),
+ MemoryLimitBytes: int64(memoryLimitBytes),
}
devices = append(devices, device)
@@ -110,6 +112,17 @@
return devices, nil
}
+// ListDevices returns the list of devices associated with a Session.
+func (s *Session) ListDevices() ([]Device, error) {
+ status := newStatus()
+ devicesList := C.TF_SessionListDevices(s.c, status.c)
+ if err := status.Err(); err != nil {
+ return nil, fmt.Errorf("SessionListDevices() failed: %v", err)
+ }
+ defer C.TF_DeleteDeviceList(devicesList)
+ return deviceSliceFromDeviceList(devicesList)
+}
+
// Run the graph with the associated session starting with the supplied feeds
// to compute the value of the requested fetches. Runs, but does not return
// Tensors for operations specified in targets.
diff --git a/tensorflow/go/session_test.go b/tensorflow/go/session_test.go
index 05ace99..c9bda00 100644
--- a/tensorflow/go/session_test.go
+++ b/tensorflow/go/session_test.go
@@ -299,3 +299,21 @@
t.Fatalf("no devices detected")
}
}
+
+func TestDeviceString(t *testing.T) {
+ d := Device{Name: "foo", Type: "bar", MemoryLimitBytes: 12345}
+ got := d.String()
+ want := "(Device: name \"foo\", type bar, memory limit 12345 bytes)"
+ if got != want {
+ t.Errorf("Got \"%s\", want \"%s\"", got, want)
+ }
+}
+
+func TestDeviceStringNoMemoryLimit(t *testing.T) {
+ d := Device{Name: "foo", Type: "bar", MemoryLimitBytes: -1}
+ got := d.String()
+ want := "(Device: name \"foo\", type bar, no memory limit)"
+ if got != want {
+ t.Errorf("Got \"%s\", want \"%s\"", got, want)
+ }
+}
diff --git a/tensorflow/lite/BUILD b/tensorflow/lite/BUILD
index bb2c53b..8fca016 100644
--- a/tensorflow/lite/BUILD
+++ b/tensorflow/lite/BUILD
@@ -1,11 +1,12 @@
-package(default_visibility = [
- "//visibility:public",
-])
+package(
+ default_visibility = ["//visibility:public"],
+)
licenses(["notice"]) # Apache 2.0
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
-load("//tensorflow/lite:build_def.bzl", "tflite_copts", "gen_selected_ops")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow:tensorflow.bzl", "if_not_windows")
exports_files(glob([
"testdata/*.bin",
@@ -35,15 +36,22 @@
visibility = ["//visibility:public"],
)
+TFLITE_DEFAULT_COPTS = if_not_windows([
+ "-Wall",
+ "-Wno-comment",
+])
+
cc_library(
name = "schema_fbs_version",
hdrs = ["version.h"],
+ copts = TFLITE_DEFAULT_COPTS,
)
cc_library(
name = "arena_planner",
srcs = ["arena_planner.cc"],
hdrs = ["arena_planner.h"],
+ copts = TFLITE_DEFAULT_COPTS,
deps = [
":graph_info",
":memory_planner",
@@ -57,12 +65,10 @@
size = "small",
srcs = ["arena_planner_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable",
],
deps = [
":arena_planner",
- "//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/lite/testing:util",
"@com_google_googletest//:gtest",
@@ -74,18 +80,21 @@
cc_library(
name = "context",
hdrs = ["context.h"],
+ copts = TFLITE_DEFAULT_COPTS,
deps = ["//tensorflow/lite/c:c_api_internal"],
)
cc_library(
name = "graph_info",
hdrs = ["graph_info.h"],
+ copts = TFLITE_DEFAULT_COPTS,
deps = ["//tensorflow/lite/c:c_api_internal"],
)
cc_library(
name = "memory_planner",
hdrs = ["memory_planner.h"],
+ copts = TFLITE_DEFAULT_COPTS,
deps = ["//tensorflow/lite/c:c_api_internal"],
)
@@ -93,6 +102,7 @@
name = "simple_memory_arena",
srcs = ["simple_memory_arena.cc"],
hdrs = ["simple_memory_arena.h"],
+ copts = TFLITE_DEFAULT_COPTS,
deps = ["//tensorflow/lite/c:c_api_internal"],
)
@@ -109,9 +119,9 @@
hdrs = [
"builtin_op_data.h",
"builtin_ops.h",
- "context.h",
"context_util.h",
],
+ deps = ["//tensorflow/lite/c:c_api_internal"],
)
exports_files(["builtin_ops.h"])
@@ -121,9 +131,7 @@
hdrs = [
"string.h",
],
- deps = [
- "//tensorflow/core:lib_platform",
- ],
+ copts = TFLITE_DEFAULT_COPTS,
)
# TODO(ahentz): investigate dependency on gemm_support requiring usage of tf_copts.
@@ -167,7 +175,7 @@
"optional_debug_tools.h",
"stderr_reporter.h",
],
- copts = tflite_copts(),
+ copts = tflite_copts() + TFLITE_DEFAULT_COPTS,
linkopts = [
] + select({
"//tensorflow:android": [
@@ -185,7 +193,7 @@
":string",
":util",
"//tensorflow/lite/c:c_api_internal",
- "//tensorflow/lite/core/api",
+ "//tensorflow/lite/core/api:api",
"//tensorflow/lite/kernels:eigen_support",
"//tensorflow/lite/kernels:gemm_support",
"//tensorflow/lite/nnapi:nnapi_lib",
@@ -203,10 +211,10 @@
name = "string_util",
srcs = ["string_util.cc"],
hdrs = ["string_util.h"],
- copts = tflite_copts(),
+ copts = TFLITE_DEFAULT_COPTS,
deps = [
- ":framework",
":string",
+ "//tensorflow/lite/c:c_api_internal",
],
)
@@ -217,6 +225,7 @@
deps = [
":framework",
":string_util",
+ "//tensorflow/lite/c:c_api_internal",
"//tensorflow/lite/testing:util",
"@com_google_googletest//:gtest",
],
@@ -246,10 +255,8 @@
name = "graph_info_test",
size = "small",
srcs = ["graph_info_test.cc"],
- tags = ["no_oss"],
deps = [
":framework",
- ":string_util",
"//tensorflow/lite/testing:util",
"@com_google_googletest//:gtest",
],
@@ -298,7 +305,12 @@
data = [
"testdata/multi_add_flex.bin",
],
- tags = ["no_windows"], # TODO(b/116667551): No weak symbols with MSVC.
+ tags = [
+ "no_gpu", # GPU + flex is not officially supported.
+ "no_windows", # TODO(b/116667551): No weak symbols with MSVC.
+ "tflite_not_portable_android",
+ "tflite_not_portable_ios",
+ ],
deps = [
":framework",
"//tensorflow/lite/core/api",
@@ -314,7 +326,6 @@
name = "mutable_op_resolver_test",
size = "small",
srcs = ["mutable_op_resolver_test.cc"],
- tags = ["no_oss"],
deps = [
":framework",
"//tensorflow/lite/testing:util",
@@ -326,7 +337,7 @@
name = "util",
srcs = ["util.cc"],
hdrs = ["util.h"],
- copts = tflite_copts(),
+ copts = TFLITE_DEFAULT_COPTS + tflite_copts(),
deps = [
"//tensorflow/lite/c:c_api_internal",
],
@@ -336,27 +347,9 @@
name = "util_test",
size = "small",
srcs = ["util_test.cc"],
- tags = ["no_oss"],
deps = [
":util",
- "//tensorflow/lite/testing:util",
+ "//tensorflow/lite/c:c_api_internal",
"@com_google_googletest//:gtest",
],
)
-
-# Test the serialization of a model with optional tensors.
-
-# Model tests
-
-#cc_library(
-# name = "models_test_utils",
-# testonly = 1,
-# hdrs = ["models/test_utils.h"],
-# deps = select({
-# "//tensorflow:android": [],
-# "//conditions:default": [
-# "@com_google_absl//absl/strings",
-# "//tensorflow/core:test",
-# ],
-# }),
-#)
diff --git a/tensorflow/lite/build_def.bzl b/tensorflow/lite/build_def.bzl
index dacb8ad..c17eddf 100644
--- a/tensorflow/lite/build_def.bzl
+++ b/tensorflow/lite/build_def.bzl
@@ -112,7 +112,8 @@
linkshared = 1,
linkstatic = 1,
testonly = 0,
- deps = []):
+ deps = [],
+ srcs = []):
"""Builds a jni binary for TFLite."""
linkopts = linkopts + [
"-Wl,--version-script", # Export only jni functions & classes.
@@ -124,6 +125,7 @@
linkshared = linkshared,
linkstatic = linkstatic,
deps = deps + [linkscript],
+ srcs = srcs,
linkopts = linkopts,
testonly = testonly,
)
@@ -457,6 +459,7 @@
native.py_test(
name = "model_coverage_test_%s_%s" % (model_name, target_op_sets.lower().replace(",", "_")),
srcs = ["model_coverage_test.py"],
+ size = "large",
main = "model_coverage_test.py",
args = [
"--model_name=%s" % model_name,
@@ -467,7 +470,6 @@
tags = [
"no_oss",
"no_windows",
- "notap",
] + tags,
deps = [
"//tensorflow/lite/testing/model_coverage:model_coverage_lib",
diff --git a/tensorflow/lite/delegates/flex/delegate_test.cc b/tensorflow/lite/delegates/flex/delegate_test.cc
index 1b2f476..ee37090 100644
--- a/tensorflow/lite/delegates/flex/delegate_test.cc
+++ b/tensorflow/lite/delegates/flex/delegate_test.cc
@@ -22,7 +22,6 @@
namespace flex {
namespace {
-using ::testing::ContainsRegex;
using ::testing::ElementsAre;
class DelegateTest : public testing::FlexModelTest {
diff --git a/tensorflow/lite/delegates/nnapi/BUILD b/tensorflow/lite/delegates/nnapi/BUILD
index c24f0f7..fd954ba 100644
--- a/tensorflow/lite/delegates/nnapi/BUILD
+++ b/tensorflow/lite/delegates/nnapi/BUILD
@@ -23,10 +23,7 @@
name = "nnapi_delegate_test",
size = "small",
srcs = ["nnapi_delegate_test.cc"],
- tags = [
- "no_oss",
- "noasan", # TODO(b/112326936): re-enable for asan once fixed.
- ],
+ tags = ["no_oss"],
deps = [
":nnapi_delegate",
"//tensorflow/lite:framework",
diff --git a/tensorflow/lite/examples/android/BUILD b/tensorflow/lite/examples/android/BUILD
index 761a603..80cefd4 100644
--- a/tensorflow/lite/examples/android/BUILD
+++ b/tensorflow/lite/examples/android/BUILD
@@ -34,7 +34,7 @@
# to reduce APK size.
assets = [
"//tensorflow/lite/examples/android/app/src/main/assets:labels_mobilenet_quant_v1_224.txt",
- "@tflite_mobilenet//:mobilenet_quant_v1_224.tflite",
+ "@tflite_mobilenet_quant//:mobilenet_v1_1.0_224_quant.tflite",
"@tflite_conv_actions_frozen//:conv_actions_frozen.tflite",
"//tensorflow/lite/examples/android/app/src/main/assets:conv_actions_labels.txt",
"@tflite_mobilenet_ssd//:mobilenet_ssd.tflite",
diff --git a/tensorflow/lite/examples/android/app/download-models.gradle b/tensorflow/lite/examples/android/app/download-models.gradle
index d2f03db..36bd177 100644
--- a/tensorflow/lite/examples/android/app/download-models.gradle
+++ b/tensorflow/lite/examples/android/app/download-models.gradle
@@ -8,13 +8,12 @@
* 3 model files will be downloaded into given folder of ext.ASSET_DIR
*/
// hard coded model files
-// LINT.IfChange
-def models = ['conv_actions_tflite.zip',
- 'mobilenet_ssd_tflite_v1.zip',
- 'mobilenet_v1_224_android_quant_2017_11_08.zip',
- 'coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip']
-// LINT.ThenChange(//tensorflow/lite/examples/android/BUILD)
+def models = ['https://storage.googleapis.com/download.tensorflow.org/models/tflite/conv_actions_tflite.zip',
+ 'https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_ssd_tflite_v1.zip',
+ 'https://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip',
+ 'http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224.tgz',
+ 'http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz']
// Root URL for model archives
def MODEL_URL = 'https://storage.googleapis.com/download.tensorflow.org/models/tflite'
@@ -30,9 +29,9 @@
import de.undercouch.gradle.tasks.download.Download
task downloadFile(type: Download){
- for (f in models) {
- def modelUrl = MODEL_URL + "/" + f
- println "Downloading ${f} from ${modelUrl}"
+ for (modelUrl in models) {
+ def localFile = modelUrl.split("/")[-1]
+ println "Downloading ${localFile} from ${modelUrl}"
src modelUrl
}
@@ -43,7 +42,12 @@
task extractModels(type: Copy) {
for (f in models) {
def localFile = f.split("/")[-1]
- from zipTree(project.ext.TMP_DIR + '/' + localFile)
+ def localExt = localFile.split("[.]")[-1]
+ if (localExt == "tgz") {
+ from tarTree(project.ext.TMP_DIR + '/' + localFile)
+ } else {
+ from zipTree(project.ext.TMP_DIR + '/' + localFile)
+ }
}
into file(project.ext.ASSET_DIR)
@@ -63,6 +67,9 @@
}
}
+
+
+
tasks.whenTaskAdded { task ->
if (task.name == 'assembleDebug') {
task.dependsOn 'extractModels'
diff --git a/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/ClassifierActivity.java b/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/ClassifierActivity.java
index dcbbefb..698251d 100644
--- a/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/ClassifierActivity.java
+++ b/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/ClassifierActivity.java
@@ -65,7 +65,7 @@
// --input_binary=true
private static final int INPUT_SIZE = 224;
- private static final String MODEL_FILE = "mobilenet_quant_v1_224.tflite";
+ private static final String MODEL_FILE = "mobilenet_v1_1.0_224_quant.tflite";
private static final String LABEL_FILE = "labels_mobilenet_quant_v1_224.txt";
private static final boolean MAINTAIN_ASPECT = true;
diff --git a/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/DetectorActivity.java b/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/DetectorActivity.java
index 87160f6b3..2feca79 100644
--- a/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/DetectorActivity.java
+++ b/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/DetectorActivity.java
@@ -52,8 +52,8 @@
private static final int TF_OD_API_INPUT_SIZE = 300;
private static final boolean TF_OD_API_IS_QUANTIZED = true;
private static final String TF_OD_API_MODEL_FILE = "detect.tflite";
- private static final String TF_OD_API_LABELS_FILE = "file:///android_asset/coco_labels_list.txt";
-
+ private static final String TF_OD_API_LABELS_FILE = "coco_labels_list.txt";
+
// Which detection model to use: by default uses Tensorflow Object Detection API frozen
// checkpoints.
private enum DetectorMode {
diff --git a/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/TFLiteObjectDetectionAPIModel.java b/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/TFLiteObjectDetectionAPIModel.java
index 9eb21de..afbf317 100644
--- a/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/TFLiteObjectDetectionAPIModel.java
+++ b/tensorflow/lite/examples/android/app/src/main/java/org/tensorflow/demo/TFLiteObjectDetectionAPIModel.java
@@ -105,8 +105,7 @@
final TFLiteObjectDetectionAPIModel d = new TFLiteObjectDetectionAPIModel();
InputStream labelsInput = null;
- String actualFilename = labelFilename.split("file:///android_asset/")[1];
- labelsInput = assetManager.open(actualFilename);
+ labelsInput = assetManager.open(labelFilename);
BufferedReader br = null;
br = new BufferedReader(new InputStreamReader(labelsInput));
String line;
diff --git a/tensorflow/lite/examples/ios/camera/CameraExampleViewController.h b/tensorflow/lite/examples/ios/camera/CameraExampleViewController.h
index fb5800e..438e6ad 100644
--- a/tensorflow/lite/examples/ios/camera/CameraExampleViewController.h
+++ b/tensorflow/lite/examples/ios/camera/CameraExampleViewController.h
@@ -17,8 +17,26 @@
#include <vector>
-#include "tensorflow/contrib/lite/kernels/register.h"
-#include "tensorflow/contrib/lite/model.h"
+// TensorFlow Lite was migrated out of `contrib/` directory. The change
+// wasn't reflected in newest CocoaPod release yet (1.12.0).
+// Change this to 0 when using a TFLite version which is newer than 1.12.0.
+// TODO(ycling): Remove the macro when we release the next version.
+#ifndef TFLITE_USE_CONTRIB_LITE
+#define TFLITE_USE_CONTRIB_LITE 1
+#endif
+
+// Set TFLITE_USE_GPU_DELEGATE to 1 to use TFLite GPU Delegate.
+// Note: TFLite GPU Delegate binary isn't releast yet, and we're working
+// on it.
+#ifndef TFLITE_USE_GPU_DELEGATE
+#define TFLITE_USE_GPU_DELEGATE 0
+#endif
+
+#if TFLITE_USE_GPU_DELEGATE && TFLITE_USE_CONTRIB_LITE
+// Sanity check.
+#error "GPU Delegate only works with newer TFLite " \
+ "after migrating out of contrib"
+#endif
@interface CameraExampleViewController
: UIViewController<UIGestureRecognizerDelegate, AVCaptureVideoDataOutputSampleBufferDelegate> {
@@ -33,10 +51,6 @@
AVCaptureSession* session;
std::vector<std::string> labels;
- std::unique_ptr<tflite::FlatBufferModel> model;
- tflite::ops::builtin::BuiltinOpResolver resolver;
- std::unique_ptr<tflite::Interpreter> interpreter;
-
double total_latency;
int total_count;
}
diff --git a/tensorflow/lite/examples/ios/camera/CameraExampleViewController.mm b/tensorflow/lite/examples/ios/camera/CameraExampleViewController.mm
index a3e6e11..48cd313 100644
--- a/tensorflow/lite/examples/ios/camera/CameraExampleViewController.mm
+++ b/tensorflow/lite/examples/ios/camera/CameraExampleViewController.mm
@@ -23,10 +23,20 @@
#include <iostream>
#include <queue>
+#if TFLITE_USE_CONTRIB_LITE
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/model.h"
#include "tensorflow/contrib/lite/op_resolver.h"
#include "tensorflow/contrib/lite/string_util.h"
+#else
+#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/op_resolver.h"
+#include "tensorflow/lite/string_util.h"
+#if TFLITE_USE_GPU_DELEGATE
+#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
+#endif
+#endif
#define LOG(x) std::cerr
@@ -34,7 +44,12 @@
// If you have your own model, modify this to the file name, and make sure
// you've added the file to your app resources too.
+#if TFLITE_USE_GPU_DELEGATE
+// GPU Delegate only supports float model now.
NSString* model_file_name = @"mobilenet_v1_1.0_224";
+#else
+NSString* model_file_name = @"mobilenet_quant_v1_224.tflite";
+#endif
NSString* model_file_type = @"tflite";
// If you have your own model, point this to the labels file.
NSString* labels_file_name = @"labels";
@@ -151,7 +166,12 @@
- (void)teardownAVCapture;
@end
-@implementation CameraExampleViewController
+@implementation CameraExampleViewController {
+ std::unique_ptr<tflite::FlatBufferModel> model;
+ tflite::ops::builtin::BuiltinOpResolver resolver;
+ std::unique_ptr<tflite::Interpreter> interpreter;
+ TfLiteDelegate* delegate;
+}
- (void)setupAVCapture {
NSError* error = nil;
@@ -363,6 +383,11 @@
}
- (void)dealloc {
+#if TFLITE_USE_GPU_DELEGATE
+ if (delegate) {
+ DeleteGpuDelegate(delegate);
+ }
+#endif
[self teardownAVCapture];
}
@@ -388,6 +413,15 @@
LoadLabels(labels_file_name, labels_file_type, &labels);
tflite::InterpreterBuilder(*model, resolver)(&interpreter);
+
+#if TFLITE_USE_GPU_DELEGATE
+ GpuDelegateOptions options;
+ options.allow_precision_loss = true;
+ options.wait_type = GpuDelegateOptions::WaitType::kActive;
+ delegate = NewGpuDelegate(&options);
+ interpreter->ModifyGraphWithDelegate(delegate);
+#endif
+
// Explicitly resize the input tensor.
{
int input = interpreter->inputs()[0];
diff --git a/tensorflow/lite/examples/ios/camera/Podfile b/tensorflow/lite/examples/ios/camera/Podfile
index 96a0d23..2e15cc6 100644
--- a/tensorflow/lite/examples/ios/camera/Podfile
+++ b/tensorflow/lite/examples/ios/camera/Podfile
@@ -1,5 +1,13 @@
platform :ios, '8.0'
inhibit_all_warnings!
+project 'tflite_camera_example.xcodeproj'
+
target 'tflite_camera_example'
- pod 'TensorFlowLite', '1.12.0'
+ # Comment 'TensorFlowLite' pod and un-comment 'TensorFlowLiteGpuExperimental'
+ # to use TFLite GPU Delegate.
+ # Note: TFLite GPU Delegate binary isn't releast yet, and we're working
+ # on it.
+
+ pod 'TensorFlowLite', '1.12.0'
+ # pod 'TensorFlowLiteGpuExperimental', '0.0.1'
diff --git a/tensorflow/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj b/tensorflow/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj
index 9522c41..9b5c2b3 100644
--- a/tensorflow/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj
+++ b/tensorflow/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj
@@ -15,6 +15,7 @@
1CDB2D4A1ED3A9CD007929E9 /* CameraExampleViewController.mm in Sources */ = {isa = PBXBuildFile; fileRef = 1CDB2D451ED3A9CD007929E9 /* CameraExampleViewController.mm */; };
54DC6C3C5F734F3A58069F0C /* libPods-tflite_camera_example.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 3BA8BF92C84895BFE59D8236 /* libPods-tflite_camera_example.a */; };
AC1F82661FBA3CBD0052BA77 /* labels.txt in Resources */ = {isa = PBXBuildFile; fileRef = AC1F82641FBA3CBD0052BA77 /* labels.txt */; };
+ AC31178921BB3FF900AFF1D2 /* mobilenet_quant_v1_224.tflite in Resources */ = {isa = PBXBuildFile; fileRef = AC31178821BB3FF900AFF1D2 /* mobilenet_quant_v1_224.tflite */; };
AC3BB41720114C400084552C /* mobilenet_v1_1.0_224.tflite in Resources */ = {isa = PBXBuildFile; fileRef = AC3BB41620114C400084552C /* mobilenet_v1_1.0_224.tflite */; };
/* End PBXBuildFile section */
@@ -36,6 +37,7 @@
3BC5BE4BBD09374D3E98F082 /* Pods-tflite_camera_example.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_camera_example.debug.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example.debug.xcconfig"; sourceTree = "<group>"; };
55ED318E8D29C8AFEF03DF1E /* Pods-tflite_camera_example.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_camera_example.release.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example.release.xcconfig"; sourceTree = "<group>"; };
AC1F82641FBA3CBD0052BA77 /* labels.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = labels.txt; sourceTree = "<group>"; };
+ AC31178821BB3FF900AFF1D2 /* mobilenet_quant_v1_224.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_quant_v1_224.tflite; sourceTree = "<group>"; };
AC3BB41620114C400084552C /* mobilenet_v1_1.0_224.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_v1_1.0_224.tflite; sourceTree = "<group>"; };
/* End PBXFileReference section */
@@ -103,6 +105,7 @@
59A3CFF31CF4E68100C4259F /* data */ = {
isa = PBXGroup;
children = (
+ AC31178821BB3FF900AFF1D2 /* mobilenet_quant_v1_224.tflite */,
AC3BB41620114C400084552C /* mobilenet_v1_1.0_224.tflite */,
AC1F82641FBA3CBD0052BA77 /* labels.txt */,
);
@@ -120,8 +123,6 @@
1C564C091ED3A92E00087306 /* Sources */,
1C564C0A1ED3A92E00087306 /* Frameworks */,
1C564C0B1ED3A92E00087306 /* Resources */,
- 00E875C3B066535AE6B77101 /* [CP] Embed Pods Frameworks */,
- 5C2D02120E3E5E09567AA946 /* [CP] Copy Pods Resources */,
);
buildRules = (
);
@@ -175,42 +176,13 @@
AC3BB41720114C400084552C /* mobilenet_v1_1.0_224.tflite in Resources */,
1C99111C1ED3B0E600A6BFB9 /* MainStoryboard_iPhone.storyboard in Resources */,
AC1F82661FBA3CBD0052BA77 /* labels.txt in Resources */,
+ AC31178921BB3FF900AFF1D2 /* mobilenet_quant_v1_224.tflite in Resources */,
);
runOnlyForDeploymentPostprocessing = 0;
};
/* End PBXResourcesBuildPhase section */
/* Begin PBXShellScriptBuildPhase section */
- 00E875C3B066535AE6B77101 /* [CP] Embed Pods Frameworks */ = {
- isa = PBXShellScriptBuildPhase;
- buildActionMask = 2147483647;
- files = (
- );
- inputPaths = (
- );
- name = "[CP] Embed Pods Frameworks";
- outputPaths = (
- );
- runOnlyForDeploymentPostprocessing = 0;
- shellPath = /bin/sh;
- shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example-frameworks.sh\"\n";
- showEnvVarsInLog = 0;
- };
- 5C2D02120E3E5E09567AA946 /* [CP] Copy Pods Resources */ = {
- isa = PBXShellScriptBuildPhase;
- buildActionMask = 2147483647;
- files = (
- );
- inputPaths = (
- );
- name = "[CP] Copy Pods Resources";
- outputPaths = (
- );
- runOnlyForDeploymentPostprocessing = 0;
- shellPath = /bin/sh;
- shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example-resources.sh\"\n";
- showEnvVarsInLog = 0;
- };
66DAEAAEE9EF6550C3A061E0 /* [CP] Check Pods Manifest.lock */ = {
isa = PBXShellScriptBuildPhase;
buildActionMask = 2147483647;
@@ -322,9 +294,7 @@
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
GCC_WARN_UNUSED_FUNCTION = YES;
GCC_WARN_UNUSED_VARIABLE = YES;
- HEADER_SEARCH_PATHS = (
- "$(inherited)",
- );
+ HEADER_SEARCH_PATHS = "$(inherited)";
IPHONEOS_DEPLOYMENT_TARGET = 8.0;
MTL_ENABLE_DEBUG_INFO = YES;
ONLY_ACTIVE_ARCH = YES;
@@ -365,9 +335,7 @@
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
GCC_WARN_UNUSED_FUNCTION = YES;
GCC_WARN_UNUSED_VARIABLE = YES;
- HEADER_SEARCH_PATHS = (
- "$(inherited)",
- );
+ HEADER_SEARCH_PATHS = "$(inherited)";
IPHONEOS_DEPLOYMENT_TARGET = 8.0;
MTL_ENABLE_DEBUG_INFO = NO;
SDKROOT = iphoneos;
diff --git a/tensorflow/lite/examples/ios/download_models.sh b/tensorflow/lite/examples/ios/download_models.sh
index ad6ccd1..4828617 100755
--- a/tensorflow/lite/examples/ios/download_models.sh
+++ b/tensorflow/lite/examples/ios/download_models.sh
@@ -53,6 +53,6 @@
file ${DOWNLOADS_DIR}/models
cp ${DOWNLOADS_DIR}/models/models/* simple/data/
-cp "${DOWNLOADS_DIR}/quantized_models/labels.txt" camera/data/
+cp ${DOWNLOADS_DIR}/models/models/* camera/data/
cp "${DOWNLOADS_DIR}/quantized_models/mobilenet_quant_v1_224.tflite" \
- 'camera/data/mobilenet_v1_1.0_224.tflite'
+ 'camera/data/mobilenet_quant_v1_224.tflite'
diff --git a/tensorflow/lite/examples/label_image/BUILD b/tensorflow/lite/examples/label_image/BUILD
index de1bfd7..4fc8648 100644
--- a/tensorflow/lite/examples/label_image/BUILD
+++ b/tensorflow/lite/examples/label_image/BUILD
@@ -63,7 +63,6 @@
data = [
"testdata/grace_hopper.bmp",
],
- tags = ["no_oss"],
deps = [
":bitmap_helpers",
"@com_google_googletest//:gtest",
diff --git a/tensorflow/lite/examples/label_image/label_image_test.cc b/tensorflow/lite/examples/label_image/label_image_test.cc
index 6b4ec2a..4db139f 100644
--- a/tensorflow/lite/examples/label_image/label_image_test.cc
+++ b/tensorflow/lite/examples/label_image/label_image_test.cc
@@ -20,8 +20,6 @@
#include "tensorflow/lite/examples/label_image/get_top_n.h"
#include "tensorflow/lite/examples/label_image/label_image.h"
-using ::testing::ElementsAreArray;
-
namespace tflite {
namespace label_image {
diff --git a/tensorflow/lite/examples/minimal/minimal.cc b/tensorflow/lite/examples/minimal/minimal.cc
index 46f8b09..9bbfee6 100644
--- a/tensorflow/lite/examples/minimal/minimal.cc
+++ b/tensorflow/lite/examples/minimal/minimal.cc
@@ -50,7 +50,7 @@
// Build the interpreter
tflite::ops::builtin::BuiltinOpResolver resolver;
- InterpreterBuilder builder(*model.get(), resolver);
+ InterpreterBuilder builder(*model, resolver);
std::unique_ptr<Interpreter> interpreter;
builder(&interpreter);
TFLITE_MINIMAL_CHECK(interpreter != nullptr);
diff --git a/tensorflow/lite/experimental/examples/lstm/unidirectional_sequence_lstm_test.py b/tensorflow/lite/experimental/examples/lstm/unidirectional_sequence_lstm_test.py
index 9c00d05..eeb48d1 100644
--- a/tensorflow/lite/experimental/examples/lstm/unidirectional_sequence_lstm_test.py
+++ b/tensorflow/lite/experimental/examples/lstm/unidirectional_sequence_lstm_test.py
@@ -111,7 +111,7 @@
# Initialize variables
init = tf.global_variables_initializer()
- self.evaluate(init)
+ sess.run(init)
for _ in range(TRAIN_STEPS):
batch_x, batch_y = self.mnist.train.next_batch(
batch_size=self.batch_size, shuffle=False)
diff --git a/tensorflow/lite/experimental/micro/examples/micro_speech/Makefile.inc b/tensorflow/lite/experimental/micro/examples/micro_speech/Makefile.inc
new file mode 100644
index 0000000..0e42329
--- /dev/null
+++ b/tensorflow/lite/experimental/micro/examples/micro_speech/Makefile.inc
@@ -0,0 +1,153 @@
+
+# Tests loading and running a speech model.
+MICRO_SPEECH_TEST_SRCS := \
+tensorflow/lite/experimental/micro/examples/micro_speech/micro_speech_test.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/tiny_conv_model_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/no_features_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/yes_features_data.cc
+ALL_SRCS += $(MICRO_SPEECH_TEST_SRCS)
+MICRO_SPEECH_TEST_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MICRO_SPEECH_TEST_SRCS))))
+MICRO_SPEECH_TEST_BINARY := $(BINDIR)micro_speech_test
+ALL_BINARIES += $(MICRO_SPEECH_TEST_BINARY)
+$(MICRO_SPEECH_TEST_BINARY): $(MICRO_SPEECH_TEST_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(MICRO_SPEECH_TEST_BINARY) $(MICRO_SPEECH_TEST_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+micro_speech_test: $(MICRO_SPEECH_TEST_BINARY)
+micro_speech_test_bin: $(MICRO_SPEECH_TEST_BINARY).bin
+test_micro_speech: $(MICRO_SPEECH_TEST_BINARY)
+ $(TEST_SCRIPT) $(MICRO_SPEECH_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
+
+# Source files that are used by multiple preprocessor tests.
+PREPROCESSOR_TEST_SHARED_SRCS := \
+tensorflow/lite/experimental/micro/examples/micro_speech/preprocessor_test.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/no_30ms_sample_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/yes_30ms_sample_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/no_power_spectrum_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/yes_power_spectrum_data.cc
+
+# Test the float reference code for feature generation.
+PREPROCESSOR_REFERENCE_TEST_SRCS = \
+$(PREPROCESSOR_TEST_SHARED_SRCS) \
+tensorflow/lite/experimental/micro/examples/micro_speech/preprocessor.cc
+ALL_SRCS += $(PREPROCESSOR_REFERENCE_TEST_SRCS)
+PREPROCESSOR_REFERENCE_TEST_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(PREPROCESSOR_REFERENCE_TEST_SRCS))))
+PREPROCESSOR_REFERENCE_TEST_BINARY := $(BINDIR)preprocessor_reference_test
+ALL_BINARIES += $(PREPROCESSOR_REFERENCE_TEST_BINARY)
+$(PREPROCESSOR_REFERENCE_TEST_BINARY): $(PREPROCESSOR_REFERENCE_TEST_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(PREPROCESSOR_REFERENCE_TEST_BINARY) $(PREPROCESSOR_REFERENCE_TEST_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+preprocessor_reference_test: $(PREPROCESSOR_REFERENCE_TEST_BINARY)
+preprocessor_reference_test_bin: $(PREPROCESSOR_REFERENCE_TEST_BINARY).bin
+test_preprocessor_reference: $(PREPROCESSOR_REFERENCE_TEST_BINARY)
+ $(TEST_SCRIPT) $(PREPROCESSOR_REFERENCE_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
+
+# Test the fixed point reference code for feature generation.
+PREPROCESSOR_FIXED_TEST_SRCS = \
+$(PREPROCESSOR_TEST_SHARED_SRCS) \
+tensorflow/lite/experimental/micro/examples/micro_speech/fixed_point/preprocessor.cc
+ALL_SRCS += $(PREPROCESSOR_FIXED_TEST_SRCS)
+PREPROCESSOR_FIXED_TEST_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(PREPROCESSOR_FIXED_TEST_SRCS))))
+PREPROCESSOR_FIXED_TEST_BINARY := $(BINDIR)preprocessor_fixed_test
+ALL_BINARIES += $(PREPROCESSOR_FIXED_TEST_BINARY)
+$(PREPROCESSOR_FIXED_TEST_BINARY): $(PREPROCESSOR_FIXED_TEST_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(PREPROCESSOR_FIXED_TEST_BINARY) $(PREPROCESSOR_FIXED_TEST_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+preprocessor_fixed_test: $(PREPROCESSOR_FIXED_TEST_BINARY)
+preprocessor_fixed_test_bin: $(PREPROCESSOR_FIXED_TEST_BINARY).bin
+test_preprocessor_fixed: $(PREPROCESSOR_FIXED_TEST_BINARY)
+ $(TEST_SCRIPT) $(PREPROCESSOR_FIXED_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
+
+# Tests the audio provider module.
+AUDIO_PROVIDER_TEST_SRCS := \
+tensorflow/lite/experimental/micro/examples/micro_speech/audio_provider_test.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/model_settings.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/audio_provider.cc
+ALL_SRCS += $(AUDIO_PROVIDER_TEST_SRCS)
+AUDIO_PROVIDER_TEST_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(AUDIO_PROVIDER_TEST_SRCS))))
+AUDIO_PROVIDER_TEST_BINARY := $(BINDIR)audio_provider_test
+ALL_BINARIES += $(AUDIO_PROVIDER_TEST_BINARY)
+$(AUDIO_PROVIDER_TEST_BINARY): $(AUDIO_PROVIDER_TEST_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(AUDIO_PROVIDER_TEST_BINARY) $(AUDIO_PROVIDER_TEST_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+audio_provider_test: $(AUDIO_PROVIDER_TEST_BINARY)
+audio_provider_test_bin: $(AUDIO_PROVIDER_TEST_BINARY).bin
+test_audio_provider: $(AUDIO_PROVIDER_TEST_BINARY)
+ $(TEST_SCRIPT) $(AUDIO_PROVIDER_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
+
+# Tests the feature provider module.
+FEATURE_PROVIDER_TEST_SRCS := \
+tensorflow/lite/experimental/micro/examples/micro_speech/feature_provider_test.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/model_settings.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/audio_provider.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/timer.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/preprocessor.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/feature_provider.cc
+ALL_SRCS += $(FEATURE_PROVIDER_TEST_SRCS)
+FEATURE_PROVIDER_TEST_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(FEATURE_PROVIDER_TEST_SRCS))))
+FEATURE_PROVIDER_TEST_BINARY := $(BINDIR)feature_provider_test
+ALL_BINARIES += $(FEATURE_PROVIDER_TEST_BINARY)
+$(FEATURE_PROVIDER_TEST_BINARY): $(FEATURE_PROVIDER_TEST_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(FEATURE_PROVIDER_TEST_BINARY) $(FEATURE_PROVIDER_TEST_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+feature_provider_test: $(FEATURE_PROVIDER_TEST_BINARY)
+feature_provider_test_bin: $(FEATURE_PROVIDER_TEST_BINARY).bin
+test_feature_provider: $(FEATURE_PROVIDER_TEST_BINARY)
+ $(TEST_SCRIPT) $(FEATURE_PROVIDER_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
+
+# Tests the timer module.
+TIMER_TEST_SRCS := \
+tensorflow/lite/experimental/micro/examples/micro_speech/timer_test.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/timer.cc
+ALL_SRCS += $(TIMER_TEST_SRCS)
+TIMER_TEST_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(TIMER_TEST_SRCS))))
+TIMER_TEST_BINARY := $(BINDIR)timer_test
+ALL_BINARIES += $(TIMER_TEST_BINARY)
+$(TIMER_TEST_BINARY): $(TIMER_TEST_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(TIMER_TEST_BINARY) $(TIMER_TEST_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+timer_test: $(TIMER_TEST_BINARY)
+timer_test_bin: $(TIMER_TEST_BINARY).bin
+test_timer: $(TIMER_TEST_BINARY)
+ $(TEST_SCRIPT) $(TIMER_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
+
+# Builds a standalone speech command recognizer binary.
+MICRO_SPEECH_SRCS := \
+tensorflow/lite/experimental/micro/examples/micro_speech/main.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/model_settings.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/audio_provider.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/feature_provider.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/timer.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/preprocessor.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/no_features_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/yes_features_data.cc \
+tensorflow/lite/experimental/micro/examples/micro_speech/tiny_conv_model_data.cc
+ALL_SRCS += $(MICRO_SPEECH_SRCS)
+MICRO_SPEECH_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MICRO_SPEECH_SRCS))))
+MICRO_SPEECH_BINARY := $(BINDIR)micro_speech
+ALL_BINARIES += $(MICRO_SPEECH_BINARY)
+$(MICRO_SPEECH_BINARY): $(MICRO_SPEECH_OBJS) $(MICROLITE_LIB_PATH)
+ @mkdir -p $(dir $@)
+ $(CXX) $(CXXFLAGS) $(INCLUDES) \
+ -o $(MICRO_SPEECH_BINARY) $(MICRO_SPEECH_OBJS) \
+ $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
+micro_speech: $(MICRO_SPEECH_BINARY)
+micro_speech_bin: $(MICRO_SPEECH_BINARY).bin
diff --git a/tensorflow/lite/experimental/micro/tools/make/Makefile b/tensorflow/lite/experimental/micro/tools/make/Makefile
index 0caf0ca..20307e2 100644
--- a/tensorflow/lite/experimental/micro/tools/make/Makefile
+++ b/tensorflow/lite/experimental/micro/tools/make/Makefile
@@ -52,29 +52,6 @@
# runtime that can be linked in to other programs.
MICROLITE_LIB_NAME := libtensorflow-microlite.a
-# Test binary for the microcontroller speech model.
-MICRO_SPEECH_TEST_SRCS := \
-tensorflow/lite/experimental/micro/examples/micro_speech/micro_speech_test.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/tiny_conv_model_data.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/no_features_data.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/yes_features_data.cc
-
-# Test binary for the microcontroller speech model.
-PREPROCESSOR_TEST_SRCS := \
-tensorflow/lite/experimental/micro/examples/micro_speech/preprocessor_test.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/no_30ms_sample_data.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/yes_30ms_sample_data.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/no_power_spectrum_data.cc \
-tensorflow/lite/experimental/micro/examples/micro_speech/yes_power_spectrum_data.cc
-
-PREPROCESSOR_REFERENCE_TEST_SRCS = \
-$(PREPROCESSOR_TEST_SRCS) \
-tensorflow/lite/experimental/micro/examples/micro_speech/preprocessor.cc
-
-PREPROCESSOR_FIXED_TEST_SRCS += \
-$(PREPROCESSOR_TEST_SRCS) \
-tensorflow/lite/experimental/micro/examples/micro_speech/fixed_point/preprocessor.cc
-
MICROLITE_TEST_SRCS := \
$(wildcard tensorflow/lite/experimental/micro/*test.cc) \
$(wildcard tensorflow/lite/experimental/micro/kernels/*test.cc)
@@ -97,9 +74,6 @@
include $(wildcard $(MAKEFILE_DIR)/targets/*_makefile.inc)
ALL_SRCS := \
- $(MICRO_SPEECH_TEST_SRCS) \
- $(PREPROCESSOR_REFERENCE_TEST_SRCS) \
- $(PREPROCESSOR_FIXED_TEST_SRCS) \
$(MICROLITE_CC_SRCS) \
$(MICROLITE_TEST_SRCS)
@@ -111,22 +85,12 @@
MICROLITE_LIB_PATH := $(LIBDIR)$(MICROLITE_LIB_NAME)
-MICRO_SPEECH_TEST_BINARY := $(BINDIR)micro_speech_test
-PREPROCESSOR_REFERENCE_TEST_BINARY := $(BINDIR)preprocessor_reference_test
-PREPROCESSOR_FIXED_TEST_BINARY := $(BINDIR)preprocessor_fixed_test
-
CXX := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}g++
CC := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}gcc
AR := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}ar
-MICRO_SPEECH_TEST_OBJS := $(addprefix $(OBJDIR), \
-$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MICRO_SPEECH_TEST_SRCS))))
-
-PREPROCESSOR_REFERENCE_TEST_OBJS := $(addprefix $(OBJDIR), \
-$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(PREPROCESSOR_REFERENCE_TEST_SRCS))))
-
-PREPROCESSOR_FIXED_TEST_OBJS := $(addprefix $(OBJDIR), \
-$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(PREPROCESSOR_FIXED_TEST_SRCS))))
+# Load the examples.
+include $(wildcard tensorflow/lite/experimental/micro/examples/*/Makefile.inc)
MICROLITE_LIB_OBJS := $(addprefix $(OBJDIR), \
$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MICROLITE_CC_SRCS))))
@@ -145,7 +109,7 @@
$(CC) $(CCFLAGS) $(INCLUDES) -c $< -o $@
# The target that's compiled if there's no command-line arguments.
-all: $(MICROLITE_LIB_PATH) $(MICRO_SPEECH_TEST_BINARY) $(PREPROCESSOR_TEST_BINARY)
+all: $(MICROLITE_LIB_PATH) $(ALL_BINARIES)
microlite: $(MICROLITE_LIB_PATH)
@@ -158,42 +122,6 @@
@mkdir -p $(dir $@)
$(AR) $(ARFLAGS) $(MICROLITE_LIB_PATH) $(MICROLITE_LIB_OBJS)
-$(MICRO_SPEECH_TEST_BINARY): $(MICRO_SPEECH_TEST_OBJS) $(MICROLITE_LIB_PATH)
- @mkdir -p $(dir $@)
- $(CXX) $(CXXFLAGS) $(INCLUDES) \
- -o $(MICRO_SPEECH_TEST_BINARY) $(MICRO_SPEECH_TEST_OBJS) \
- $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
-
-micro_speech_test: $(MICRO_SPEECH_TEST_BINARY)
-micro_speech_test_bin: $(MICRO_SPEECH_TEST_BINARY).bin
-
-test_micro_speech: $(MICRO_SPEECH_TEST_BINARY)
- $(TEST_SCRIPT) $(MICRO_SPEECH_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
-
-$(PREPROCESSOR_REFERENCE_TEST_BINARY): $(PREPROCESSOR_REFERENCE_TEST_OBJS) $(MICROLITE_LIB_PATH)
- @mkdir -p $(dir $@)
- $(CXX) $(CXXFLAGS) $(INCLUDES) \
- -o $(PREPROCESSOR_REFERENCE_TEST_BINARY) $(PREPROCESSOR_REFERENCE_TEST_OBJS) \
- $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
-
-preprocessor_reference_test: $(PREPROCESSOR_REFERENCE_TEST_BINARY)
-preprocessor_reference_test_bin: $(PREPROCESSOR_REFERENCE_TEST_BINARY).bin
-
-test_preprocessor_reference: $(PREPROCESSOR_REFERENCE_TEST_BINARY)
- $(TEST_SCRIPT) $(PREPROCESSOR_REFERENCE_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
-
-$(PREPROCESSOR_FIXED_TEST_BINARY): $(PREPROCESSOR_FIXED_TEST_OBJS) $(MICROLITE_LIB_PATH)
- @mkdir -p $(dir $@)
- $(CXX) $(CXXFLAGS) $(INCLUDES) \
- -o $(PREPROCESSOR_FIXED_TEST_BINARY) $(PREPROCESSOR_FIXED_TEST_OBJS) \
- $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS)
-
-preprocessor_fixed_test: $(PREPROCESSOR_FIXED_TEST_BINARY)
-preprocessor_fixed_test_bin: $(PREPROCESSOR_FIXED_TEST_BINARY).bin
-
-test_preprocessor_fixed: $(PREPROCESSOR_FIXED_TEST_BINARY)
- $(TEST_SCRIPT) $(PREPROCESSOR_FIXED_TEST_BINARY) '~~~ALL TESTS PASSED~~~'
-
$(BINDIR)%_test : $(OBJDIR)%_test.o $(MICROLITE_LIB_PATH)
@mkdir -p $(dir $@)
$(CXX) $(CXXFLAGS) $(INCLUDES) \
@@ -203,8 +131,6 @@
$(BINDIR)%.test_target: $(BINDIR)%_test
$(TEST_SCRIPT) $< '~~~ALL TESTS PASSED~~~'
-$(info $(MICROLITE_TEST_TARGETS))
-
test: test_micro_speech $(MICROLITE_TEST_TARGETS)
# Gets rid of all generated files.
diff --git a/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc b/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc
index 4367fe7..84ab164 100644
--- a/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc
+++ b/tensorflow/lite/experimental/microfrontend/audio_microfrontend.cc
@@ -142,7 +142,8 @@
if (output.values != nullptr) {
frame_buffer[frame_index].reserve(output.size);
- for (int i = 0; i < output.size; ++i) {
+ int i;
+ for (i = 0; i < output.size; ++i) {
frame_buffer[frame_index].push_back(static_cast<T>(output.values[i]) /
data->out_scale);
}
@@ -152,9 +153,10 @@
int index = 0;
std::vector<T> pad(data->state->filterbank.num_channels, 0);
- for (int anchor = 0; anchor < frame_buffer.size();
- anchor += data->frame_stride) {
- for (int frame = anchor - data->left_context;
+ int anchor;
+ for (anchor = 0; anchor < frame_buffer.size(); anchor += data->frame_stride) {
+ int frame;
+ for (frame = anchor - data->left_context;
frame <= anchor + data->right_context; ++frame) {
std::vector<T>* feature;
if (data->zero_padding && (frame < 0 || frame >= frame_buffer.size())) {
diff --git a/tensorflow/lite/experimental/microfrontend/audio_microfrontend_test.cc b/tensorflow/lite/experimental/microfrontend/audio_microfrontend_test.cc
index a9119d0..e3a0e06 100644
--- a/tensorflow/lite/experimental/microfrontend/audio_microfrontend_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/audio_microfrontend_test.cc
@@ -140,13 +140,16 @@
// Mimic padding behaviour with zero_padding = true.
std::vector<int> output_flattened;
- for (int anchor = 0; anchor < output.size();
+ int anchor;
+ for (anchor = 0; anchor < output.size();
anchor += micro_frontend->num_frame_stride()) {
- for (int frame = anchor - micro_frontend->num_left_context();
+ int frame;
+ for (frame = anchor - micro_frontend->num_left_context();
frame <= anchor + micro_frontend->num_right_context(); ++frame) {
if (frame < 0 || frame >= output.size()) {
// Padding with zeros.
- for (int j = 0; j < num_frequency_per_frame; ++j) {
+ int j;
+ for (j = 0; j < num_frequency_per_frame; ++j) {
output_flattened.push_back(0.0);
}
} else {
diff --git a/tensorflow/lite/experimental/microfrontend/lib/fft_test.cc b/tensorflow/lite/experimental/microfrontend/lib/fft_test.cc
index 7c1ee2d..1b754c1 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/fft_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/fft_test.cc
@@ -38,7 +38,8 @@
{-887, 0}, {3000, 3000}, {0, -6401}, {-3000, 3000}, {886, 0}, {118, 119},
{0, 25}, {9, -10}, {19, 0}, {9, 9}, {0, 0}};
ASSERT_EQ(state.fft_size / 2 + 1, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i <= state.fft_size / 2; ++i) {
+ int i;
+ for (i = 0; i <= state.fft_size / 2; ++i) {
EXPECT_EQ(state.output[i].real, expected[i].real);
EXPECT_EQ(state.output[i].imag, expected[i].imag);
}
diff --git a/tensorflow/lite/experimental/microfrontend/lib/filterbank_io.c b/tensorflow/lite/experimental/microfrontend/lib/filterbank_io.c
index 2dbb4b3..6ce4c7c 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/filterbank_io.c
+++ b/tensorflow/lite/experimental/microfrontend/lib/filterbank_io.c
@@ -17,7 +17,8 @@
static void PrintArray(FILE* fp, const char* name, const int16_t* values,
size_t size) {
fprintf(fp, "static int16_t filterbank_%s[] = {", name);
- for (int i = 0; i < size; ++i) {
+ int i;
+ for (i = 0; i < size; ++i) {
fprintf(fp, "%d", values[i]);
if (i < size - 1) {
fprintf(fp, ", ");
diff --git a/tensorflow/lite/experimental/microfrontend/lib/filterbank_test.cc b/tensorflow/lite/experimental/microfrontend/lib/filterbank_test.cc
index 808d527..41f0064 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/filterbank_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/filterbank_test.cc
@@ -71,7 +71,8 @@
const int16_t expected[] = {0, 4, 8};
ASSERT_EQ(state.num_channels + 1, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i <= state.num_channels; ++i) {
+ int i;
+ for (i = 0; i <= state.num_channels; ++i) {
EXPECT_EQ(state.channel_frequency_starts[i], expected[i]);
}
@@ -85,7 +86,8 @@
const int16_t expected[] = {0, 8, 16};
ASSERT_EQ(state.num_channels + 1, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i <= state.num_channels; ++i) {
+ int i;
+ for (i = 0; i <= state.num_channels; ++i) {
EXPECT_EQ(state.channel_weight_starts[i], expected[i]);
}
@@ -99,7 +101,8 @@
const int16_t expected[] = {8, 8, 8};
ASSERT_EQ(state.num_channels + 1, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i <= state.num_channels; ++i) {
+ int i;
+ for (i = 0; i <= state.num_channels; ++i) {
EXPECT_EQ(state.channel_widths[i], expected[i]);
}
@@ -117,7 +120,8 @@
ASSERT_EQ(state.channel_weight_starts[state.num_channels] +
state.channel_widths[state.num_channels],
sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
+ int i;
+ for (i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
EXPECT_EQ(state.weights[i], expected[i]);
}
@@ -135,7 +139,8 @@
ASSERT_EQ(state.channel_weight_starts[state.num_channels] +
state.channel_widths[state.num_channels],
sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
+ int i;
+ for (i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
EXPECT_EQ(state.unweights[i], expected[i]);
}
@@ -154,7 +159,8 @@
int32_t* energy = reinterpret_cast<int32_t*>(fake_fft);
FilterbankConvertFftComplexToEnergy(&state, fake_fft, energy);
- for (int i = state.start_index; i < state.end_index; ++i) {
+ int i;
+ for (i = state.start_index; i < state.end_index; ++i) {
EXPECT_EQ(energy[i], kEnergy[i]);
}
}
@@ -167,7 +173,8 @@
FilterbankAccumulateChannels(&state, kEnergy);
ASSERT_EQ(state.num_channels + 1, sizeof(kWork) / sizeof(kWork[0]));
- for (int i = 0; i <= state.num_channels; ++i) {
+ int i;
+ for (i = 0; i <= state.num_channels; ++i) {
EXPECT_EQ(state.work[i], kWork[i]);
}
@@ -184,7 +191,8 @@
const uint32_t expected[] = {247311, 508620};
ASSERT_EQ(state.num_channels, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.num_channels; ++i) {
+ int i;
+ for (i = 0; i < state.num_channels; ++i) {
EXPECT_EQ(scaled_filterbank[i], expected[i]);
}
diff --git a/tensorflow/lite/experimental/microfrontend/lib/frontend_test.cc b/tensorflow/lite/experimental/microfrontend/lib/frontend_test.cc
index 993e866..a6faa1f 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/frontend_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/frontend_test.cc
@@ -64,7 +64,8 @@
const uint16_t expected[] = {479, 425};
ASSERT_EQ(output.size, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < output.size; ++i) {
+ int i;
+ for (i = 0; i < output.size; ++i) {
EXPECT_EQ(output.values[i], expected[i]);
}
@@ -86,7 +87,8 @@
const int16_t expected[] = {436, 378};
ASSERT_EQ(output.size, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < output.size; ++i) {
+ int i;
+ for (i = 0; i < output.size; ++i) {
EXPECT_EQ(output.values[i], expected[i]);
}
diff --git a/tensorflow/lite/experimental/microfrontend/lib/log_scale.c b/tensorflow/lite/experimental/microfrontend/lib/log_scale.c
index 54f370e..149ec7c 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/log_scale.c
+++ b/tensorflow/lite/experimental/microfrontend/lib/log_scale.c
@@ -63,7 +63,8 @@
const int scale_shift = state->scale_shift;
uint16_t* output = (uint16_t*) signal;
uint16_t* ret = output;
- for (int i = 0; i < signal_size; ++i) {
+ int i;
+ for (i = 0; i < signal_size; ++i) {
uint32_t value = *signal++;
if (state->enable_log) {
if (correction_bits < 0) {
diff --git a/tensorflow/lite/experimental/microfrontend/lib/log_scale_test.cc b/tensorflow/lite/experimental/microfrontend/lib/log_scale_test.cc
index 91ca657..1ea0842 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/log_scale_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/log_scale_test.cc
@@ -34,7 +34,8 @@
kCorrectionBits);
const uint16_t expected[] = {479, 425};
- for (int i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
+ int i;
+ for (i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
EXPECT_EQ(output[i], expected[i]);
}
}
@@ -50,7 +51,8 @@
kCorrectionBits);
const uint16_t expected[] = {65535, 45998};
- for (int i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
+ int i;
+ for (i = 0; i < sizeof(expected) / sizeof(expected[0]); ++i) {
EXPECT_EQ(output[i], expected[i]);
}
}
diff --git a/tensorflow/lite/experimental/microfrontend/lib/noise_reduction_test.cc b/tensorflow/lite/experimental/microfrontend/lib/noise_reduction_test.cc
index 1614056..13d58b2 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/noise_reduction_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/noise_reduction_test.cc
@@ -44,7 +44,8 @@
const uint32_t expected[] = {6321887, 31248341};
ASSERT_EQ(state.num_channels, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.num_channels; ++i) {
+ int i;
+ for (i = 0; i < state.num_channels; ++i) {
EXPECT_EQ(state.estimate[i], expected[i]);
}
@@ -60,7 +61,8 @@
const uint32_t expected[] = {241137, 478104};
ASSERT_EQ(state.num_channels, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.num_channels; ++i) {
+ int i;
+ for (i = 0; i < state.num_channels; ++i) {
EXPECT_EQ(signal[i], expected[i]);
}
diff --git a/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control.c b/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control.c
index b49eb30..8ccc2fd 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control.c
+++ b/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control.c
@@ -47,7 +47,8 @@
void PcanGainControlApply(struct PcanGainControlState* state,
uint32_t* signal) {
- for (int i = 0; i < state->num_channels; ++i) {
+ int i;
+ for (i = 0; i < state->num_channels; ++i) {
const uint32_t gain = WideDynamicFunction(state->noise_estimate[i],
state->gain_lut);
const uint32_t snr = ((uint64_t) signal[i] * gain) >> state->snr_shift;
diff --git a/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_test.cc b/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_test.cc
index 830db89..7c92d2d 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_test.cc
@@ -49,7 +49,8 @@
const uint32_t expected[] = {3578, 1533};
ASSERT_EQ(state.num_channels, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.num_channels; ++i) {
+ int i;
+ for (i = 0; i < state.num_channels; ++i) {
EXPECT_EQ(signal[i], expected[i]);
}
diff --git a/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_util.c b/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_util.c
index dbe44c4..5201cf0 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_util.c
+++ b/tensorflow/lite/experimental/microfrontend/lib/pcan_gain_control_util.c
@@ -62,7 +62,8 @@
state->gain_lut[0] = PcanGainLookupFunction(config, input_bits, 0);
state->gain_lut[1] = PcanGainLookupFunction(config, input_bits, 1);
state->gain_lut -= 6;
- for (int interval = 2; interval <= kWideDynamicFunctionBits; ++interval) {
+ int interval;
+ for (interval = 2; interval <= kWideDynamicFunctionBits; ++interval) {
const uint32_t x0 = (uint32_t) 1 << (interval - 1);
const uint32_t x1 = x0 + (x0 >> 1);
const uint32_t x2 = (interval == kWideDynamicFunctionBits)
diff --git a/tensorflow/lite/experimental/microfrontend/lib/window_io.c b/tensorflow/lite/experimental/microfrontend/lib/window_io.c
index ed4ac5e..d12cac2 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/window_io.c
+++ b/tensorflow/lite/experimental/microfrontend/lib/window_io.c
@@ -16,7 +16,8 @@
void WindowWriteMemmapPreamble(FILE* fp, const struct WindowState* state) {
fprintf(fp, "static int16_t window_coefficients[] = {\n");
- for (int i = 0; i < state->size; ++i) {
+ int i;
+ for (i = 0; i < state->size; ++i) {
fprintf(fp, "%d", state->coefficients[i]);
if (i < state->size - 1) {
fprintf(fp, ", ");
diff --git a/tensorflow/lite/experimental/microfrontend/lib/window_test.cc b/tensorflow/lite/experimental/microfrontend/lib/window_test.cc
index 8c6c191..60f1144 100644
--- a/tensorflow/lite/experimental/microfrontend/lib/window_test.cc
+++ b/tensorflow/lite/experimental/microfrontend/lib/window_test.cc
@@ -48,7 +48,8 @@
3843, 3541, 3145, 2681, 2177, 1664, 1176,
743, 391, 144, 16};
ASSERT_EQ(state.size, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.size; ++i) {
+ int i;
+ for (i = 0; i < state.size; ++i) {
EXPECT_EQ(state.coefficients[i], expected[i]);
}
@@ -64,7 +65,8 @@
&state, kFakeAudioData,
sizeof(kFakeAudioData) / sizeof(kFakeAudioData[0]), &num_samples_read));
- for (int i = kStepSamples; i < kWindowSamples; ++i) {
+ int i;
+ for (i = kStepSamples; i < kWindowSamples; ++i) {
EXPECT_EQ(state.input[i - kStepSamples], kFakeAudioData[i]);
}
@@ -84,7 +86,8 @@
0, 1151, 0, -5944, 0, 13311, 0, -21448, 0, 28327, 0, -32256, 0, 32255,
0, -28328, 0, 21447, 0, -13312, 0, 5943, 0, -1152, 0};
ASSERT_EQ(state.size, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.size; ++i) {
+ int i;
+ for (i = 0; i < state.size; ++i) {
EXPECT_EQ(state.output[i], expected[i]);
}
@@ -122,7 +125,8 @@
0, -1152, 0, 5943, 0, -13312, 0, 21447, 0, -28328, 0, 32255, 0, -32256,
0, 28327, 0, -21448, 0, 13311, 0, -5944, 0, 1151, 0};
ASSERT_EQ(state.size, sizeof(expected) / sizeof(expected[0]));
- for (int i = 0; i < state.size; ++i) {
+ int i;
+ for (i = 0; i < state.size; ++i) {
EXPECT_EQ(state.output[i], expected[i]);
}
diff --git a/tensorflow/lite/experimental/microfrontend/ops/audio_microfrontend_op.cc b/tensorflow/lite/experimental/microfrontend/ops/audio_microfrontend_op.cc
index 51094a9..9f2ea7e 100644
--- a/tensorflow/lite/experimental/microfrontend/ops/audio_microfrontend_op.cc
+++ b/tensorflow/lite/experimental/microfrontend/ops/audio_microfrontend_op.cc
@@ -250,7 +250,8 @@
if (output.values != nullptr) {
frame_buffer[frame_index].reserve(output.size);
- for (int i = 0; i < output.size; ++i) {
+ int i;
+ for (i = 0; i < output.size; ++i) {
frame_buffer[frame_index].push_back(static_cast<T>(output.values[i]) /
out_scale_);
}
@@ -261,9 +262,10 @@
int index = 0;
std::vector<T> pad(config_.filterbank.num_channels, 0);
- for (int anchor = 0; anchor < frame_buffer.size();
- anchor += frame_stride_) {
- for (int frame = anchor - left_context_; frame <= anchor + right_context_;
+ int anchor;
+ for (anchor = 0; anchor < frame_buffer.size(); anchor += frame_stride_) {
+ int frame;
+ for (frame = anchor - left_context_; frame <= anchor + right_context_;
++frame) {
std::vector<T>* feature;
if (zero_padding_ && (frame < 0 || frame >= frame_buffer.size())) {
diff --git a/tensorflow/lite/experimental/microfrontend/python/kernel_tests/audio_microfrontend_op_test.py b/tensorflow/lite/experimental/microfrontend/python/kernel_tests/audio_microfrontend_op_test.py
index 561f5f7..3ce8617 100644
--- a/tensorflow/lite/experimental/microfrontend/python/kernel_tests/audio_microfrontend_op_test.py
+++ b/tensorflow/lite/experimental/microfrontend/python/kernel_tests/audio_microfrontend_op_test.py
@@ -21,6 +21,7 @@
import tensorflow as tf
from tensorflow.lite.experimental.microfrontend.python.ops import audio_microfrontend_op as frontend_op
+from tensorflow.python.framework import test_util
SAMPLE_RATE = 1000
WINDOW_SIZE = 25
@@ -33,6 +34,7 @@
class AudioFeatureGenerationTest(tf.test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testSimple(self):
with self.test_session():
audio = tf.constant(
@@ -51,6 +53,7 @@
self.assertAllEqual(filterbanks.eval(),
[[479, 425], [436, 378], [410, 350], [391, 325]])
+ @test_util.run_v1_only("b/120545219")
def testSimpleFloatScaled(self):
with self.test_session():
audio = tf.constant(
@@ -72,6 +75,7 @@
[[7.484375, 6.640625], [6.8125, 5.90625],
[6.40625, 5.46875], [6.109375, 5.078125]])
+ @test_util.run_v1_only("b/120545219")
def testStacking(self):
with self.test_session():
audio = tf.constant(
@@ -114,6 +118,7 @@
[[479, 425, 479, 425, 436, 378], [479, 425, 436, 378, 410, 350],
[436, 378, 410, 350, 391, 325], [410, 350, 391, 325, 391, 325]])
+ @test_util.run_v1_only("b/120545219")
def testStackingDropFrame(self):
with self.test_session():
audio = tf.constant(
diff --git a/tensorflow/lite/g3doc/_book.yaml b/tensorflow/lite/g3doc/_book.yaml
index a51c7a6..0c79e79 100644
--- a/tensorflow/lite/g3doc/_book.yaml
+++ b/tensorflow/lite/g3doc/_book.yaml
@@ -77,8 +77,36 @@
- title: Optimizing for mobile
path: /lite/tfmobile/optimizing
+ # - name: Models
+ # contents:
+ # - title: Overview
+ # path: /lite/models/
+ # - heading: Beginner
+ # style: divider
+ # - title: Image labeling
+ # section:
+ # - title: Overview
+ # path: /lite/models/image/label/overview
+ # - title: Android
+ # path: /lite/models/image/label/android
+ # - title: iOS
+ # path: /lite/models/image/label/ios
+ # - heading: Advanced
+ # style: divider
+ # - heading: Image
+ # - title: Image classification
+ # path: /lite/models/image/classification/
+ # - heading: Audio
+ # - title: Hot word detection
+ # path: /lite/models/audio/hot_word/
+ # - heading: Text
+ # - title: Text classification
+ # path: /lite/models/text/classification/
+
- name: API
skip_translation: true
contents:
- title: API
path: /api_docs/python/tf/lite
+
+- include: /_upper_tabs_right.yaml
diff --git a/tensorflow/lite/g3doc/apis.md b/tensorflow/lite/g3doc/apis.md
index e9fa24b..b15159c 100644
--- a/tensorflow/lite/g3doc/apis.md
+++ b/tensorflow/lite/g3doc/apis.md
@@ -304,6 +304,13 @@
* `long`
* `byte`
+`String` types are also supported, but they are encoded differently than the
+primitive types. In particular, the shape of a string Tensor dictates the number
+and arrangement of strings in the Tensor, with each element itself being a
+variable length string. In this sense, the (byte) size of the Tensor cannot be
+computed from the shape and type alone, and consequently strings cannot be
+provided as a single, flat `ByteBuffer` argument.
+
If other data types, including boxed types like `Integer` and `Float`, are used,
an `IllegalArgumentException` will be thrown.
@@ -345,13 +352,12 @@
```
where each entry in `inputs` corresponds to an input tensor and
-`map_of_indices_to_outputs` maps indices of output tensors to the
-corresponding output data. In both cases the tensor indices should correspond to
-the values given to the [TensorFlow Lite Optimized Converter](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/toco/g3doc/cmdline_examples.md)
+`map_of_indices_to_outputs` maps indices of output tensors to the corresponding
+output data. In both cases the tensor indices should correspond to the values
+given to the [TensorFlow Lite Optimized Converter](convert/cmdline_examples.md)
when the model was created. Be aware that the order of tensors in `input` must
match the order given to the `TensorFlow Lite Optimized Converter`.
-
The Java API also provides convenient functions for app developers to get the
index of any model input or output using a tensor name:
diff --git a/tensorflow/lite/g3doc/convert/cmdline_examples.md b/tensorflow/lite/g3doc/convert/cmdline_examples.md
index 59f26b3..169f2d9 100644
--- a/tensorflow/lite/g3doc/convert/cmdline_examples.md
+++ b/tensorflow/lite/g3doc/convert/cmdline_examples.md
@@ -94,9 +94,10 @@
### Convert a TensorFlow GraphDef for quantized inference <a name="graphdef_quant"></a>
The TensorFlow Lite Converter is compatible with fixed point quantization models
-described [here](https://www.tensorflow.org/performance/quantization). These are
-float models with `FakeQuant*` ops inserted at the boundaries of fused layers
-to record min-max range information. This generates a quantized inference
+described
+[here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/quantize/README.md).
+These are float models with `FakeQuant*` ops inserted at the boundaries of fused
+layers to record min-max range information. This generates a quantized inference
workload that reproduces the quantization behavior that was used during
training.
diff --git a/tensorflow/lite/g3doc/convert/python_api.md b/tensorflow/lite/g3doc/convert/python_api.md
index b914a34..4d2c736 100644
--- a/tensorflow/lite/g3doc/convert/python_api.md
+++ b/tensorflow/lite/g3doc/convert/python_api.md
@@ -19,9 +19,9 @@
## API
-The API for converting TensorFlow models to TensorFlow Lite as of TensorFlow 1.9
-is `tf.lite.TFLiteConverter`. The API for calling the Python intepreter
-is `tf.lite.Interpreter`.
+The API for converting TensorFlow models to TensorFlow Lite is
+`tf.lite.TFLiteConverter`. The API for calling the Python interpreter is
+`tf.lite.Interpreter`.
`TFLiteConverter` provides class methods based on the original format of the
model. `TFLiteConverter.from_session()` is available for GraphDefs.
diff --git a/tensorflow/lite/g3doc/models/_index.yaml b/tensorflow/lite/g3doc/models/_index.yaml
new file mode 100644
index 0000000..f4d8bc4
--- /dev/null
+++ b/tensorflow/lite/g3doc/models/_index.yaml
@@ -0,0 +1,125 @@
+project_path: /lite/_project.yaml
+book_path: /lite/_book.yaml
+description: <!--no description-->
+landing_page:
+ body_class: tfo-hide-page-nav
+ custom_css_path: /site-assets/css/style.css
+ show_side_navs: true
+ rows:
+
+ # Hero
+ - classname: >
+ devsite-landing-row-50
+ devsite-landing-row-large-headings
+ devsite-landing-row-no-image-background
+ foreground: theme
+ items:
+ - heading: Models marketplace
+ description: >
+ The TensorFlow Lite models marketplace, your neighborhood model shoppe.
+ image_path: /resources/images/tflite-card-16x9.png
+
+ # Features
+ - background: grey
+ items:
+ - heading: Optimized for mobile
+ description: >
+ Machine learning can make your apps more engaging, personalized, and
+ helpful, and provides solutions that are optimized to run on-device.
+ - heading: Built with Google expertise
+ description: >
+ Models offer the technologies that have long powered Google's own
+ experiences on mobile.
+ - heading: Approachable and comprehensive
+ description: >
+ Use out-of-the-box solutions (base APIs) or custom models, running
+ on-device or in the Cloud, depending on your specific needs.
+
+ # Beginner models
+ - classname: devsite-landing-row-100
+ heading: "Build machine learning into your apps"
+ items:
+ - heading: >
+ Image labeling
+ description: >
+ Identify objects, locations, activities, animal species, products, and
+ more
+ icon:
+ path: ../images/landing-page/assistant_logo.png
+ path: /lite/image/labeling/
+ - heading: >
+ Text recognition (OCR)
+ description: >
+ Recognize and extract text from images
+ icon:
+ path: ../images/landing-page/assistant_logo.png
+ path: /lite/image/labeling/
+ - heading: >
+ Face detection
+ description: >
+ Detect faces and facial landmarks
+ icon:
+ path: ../images/landing-page/assistant_logo.png
+ path: /lite/image/labeling/
+
+ - items:
+ - heading: >
+ Barcode scanning
+ description: >
+ Scan and process barcodes
+ icon:
+ path: ../images/landing-page/assistant_logo.png
+ path: /lite/image/labeling/
+ - heading: >
+ Landmark detection
+ description: >
+ Identify popular landmarks in an image
+ icon:
+ path: ../images/landing-page/assistant_logo.png
+ path: /lite/image/labeling/
+ - heading: >
+ Smart reply
+ description: >
+ Provide suggested text snippet that fits context
+ icon:
+ path: ../images/landing-page/assistant_logo.png
+ path: /lite/image/labeling/
+
+ # Custom models
+ - classname: >
+ devsite-landing-row-no-image-background
+ devsite-landing-row-50
+ devsite-landing-row-large-headings
+ foreground: theme
+ background: grey
+ items:
+ - heading: Custom models
+ description: >
+ <p>If models don’t cover your use cases, you can always
+ bring your own existing TensorFlow Lite models. Just upload your model,
+ and we’ll take care of hosting and serving it to your app.</p>
+
+ <p>Models acts as an API layer to your custom model, making it easy to
+ run and use. In addition to deploying your models, we are releasing an
+ experimental model compression flow that aims to reduce model size (up
+ to orders of magnitudes) while maintaining similar accuracy. Sign up at
+ <a href="https://g.co/firebase/signup">g.co/firebase/signup</a></p>
+
+ <p>And if you’re new to machine learning and want more information on
+ custom models for mobile, you can <a
+ href="//www.tensorflow.org/lite/">learn more about TensorFlow
+ Lite.</a></p>
+ image_path: /resources/images/tflite-card-16x9.png
+ image_left: true
+ - classname: devsite-landing-row-large-headings
+ foreground: theme
+ items:
+ - heading: Just the beginning
+ description: >
+ Our ultimate goal is to reduce idea–to–implementation cycles and make AI
+ an essential and intuitive part of a developer's toolkit. We will do so
+ by continuing to add new Base APIs that leverage Google’s machine
+ learning expertise. Base APIs will ultimately cover significantly more
+ use cases in the vision, speech, and text fields. We will also continue
+ to simplify use of custom models, adding tools to deploy, compress, and
+ create them.
diff --git a/tensorflow/lite/g3doc/models/image/label/android.md b/tensorflow/lite/g3doc/models/image/label/android.md
new file mode 100644
index 0000000..9cd54aa
--- /dev/null
+++ b/tensorflow/lite/g3doc/models/image/label/android.md
@@ -0,0 +1,3 @@
+# Android
+
+lorem
diff --git a/tensorflow/lite/g3doc/models/image/label/ios.md b/tensorflow/lite/g3doc/models/image/label/ios.md
new file mode 100644
index 0000000..904c645
--- /dev/null
+++ b/tensorflow/lite/g3doc/models/image/label/ios.md
@@ -0,0 +1,3 @@
+# iOS
+
+lorem
diff --git a/tensorflow/lite/g3doc/models/image/label/overview.md b/tensorflow/lite/g3doc/models/image/label/overview.md
new file mode 100644
index 0000000..b3d9133
--- /dev/null
+++ b/tensorflow/lite/g3doc/models/image/label/overview.md
@@ -0,0 +1,8 @@
+# Overview
+
+Image labeling gives you insight into the content of images. When you use the
+API, you get a list of the entities that were recognized: people, things,
+places, activities, and so on. Each label found comes with a score that
+indicates the confidence the ML model has in its relevance. With this
+information, you can perform tasks such as automatic metadata generation
+and content moderation.
diff --git a/tensorflow/lite/java/demo/app/build.gradle b/tensorflow/lite/java/demo/app/build.gradle
index 05301eb..b8fc282 100644
--- a/tensorflow/lite/java/demo/app/build.gradle
+++ b/tensorflow/lite/java/demo/app/build.gradle
@@ -40,6 +40,15 @@
url 'https://google.bintray.com/tensorflow'
}
}
+allprojects {
+ repositories {
+ // Uncomment if you want to use a local repo.
+ // mavenLocal()
+ jcenter()
+ }
+}
+
+
dependencies {
compile fileTree(dir: 'libs', include: ['*.jar'])
@@ -49,31 +58,66 @@
compile 'com.android.support:support-annotations:25.3.1'
compile 'com.android.support:support-v13:25.2.0'
+ // Build off of nightly TensorFlow Lite
compile 'org.tensorflow:tensorflow-lite:0.0.0-nightly'
+ // Use local TensorFlow library
+ // compile 'org.tensorflow:tensorflow-lite-local:0.0.0'
}
-def modelDownloadUrl = "https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip"
-def localCache = "build/intermediates/mobilenet_v1_224_android_quant_2017_11_08.zip"
def targetFolder = "src/main/assets"
+def modelFloatDownloadUrl = "http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224.tgz"
+def modelQuantDownloadUrl = "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz"
+def localCacheFloat = "build/intermediates/mobilenet_v1_1.0_224.tgz"
+def localCacheQuant = "build/intermediates/mmobilenet_v1_1.0_224_quant.tgz"
-task downloadModel(type: DownloadUrlTask) {
+
+task downloadModelFloat(type: DownloadUrlTask) {
doFirst {
- println "Downloading ${modelDownloadUrl}"
+ println "Downloading ${modelFloatDownloadUrl}"
}
- sourceUrl = "${modelDownloadUrl}"
- target = file("${localCache}")
+ sourceUrl = "${modelFloatDownloadUrl}"
+ target = file("${localCacheFloat}")
}
-task unzipModel(type: Copy, dependsOn: 'downloadModel') {
+task downloadModelQuant(type: DownloadUrlTask) {
doFirst {
- println "Unzipping ${localCache}"
+ println "Downloading ${modelQuantDownloadUrl}"
}
- from zipTree("${localCache}")
+ sourceUrl = "${modelQuantDownloadUrl}"
+ target = file("${localCacheQuant}")
+}
+
+task unzipModelFloat(type: Copy, dependsOn: 'downloadModelFloat') {
+ doFirst {
+ println "Unzipping ${localCacheFloat}"
+ }
+ from tarTree("${localCacheFloat}")
into "${targetFolder}"
}
+task unzipModelQuant(type: Copy, dependsOn: 'downloadModelQuant') {
+ doFirst {
+ println "Unzipping ${localCacheQuant}"
+ }
+ from tarTree("${localCacheQuant}")
+ into "${targetFolder}"
+}
+
+task cleanUnusedFiles(type: Delete, dependsOn: ['unzipModelFloat', 'unzipModelQuant']) {
+ delete fileTree("${targetFolder}").matching {
+ include "*.pb"
+ include "*.ckpt.*"
+ include "*.pbtxt.*"
+ include "*.quant_info.*"
+ include "*.meta"
+ }
+}
+
+
// Ensure the model file is downloaded and extracted before every build
-preBuild.dependsOn unzipModel
+preBuild.dependsOn unzipModelFloat
+preBuild.dependsOn unzipModelQuant
+preBuild.dependsOn cleanUnusedFiles
class DownloadUrlTask extends DefaultTask {
@Input
@@ -87,3 +131,4 @@
ant.get(src: sourceUrl, dest: target)
}
}
+
diff --git a/tensorflow/lite/java/demo/app/src/main/BUILD b/tensorflow/lite/java/demo/app/src/main/BUILD
index df8a024..9a7c1d0 100644
--- a/tensorflow/lite/java/demo/app/src/main/BUILD
+++ b/tensorflow/lite/java/demo/app/src/main/BUILD
@@ -10,7 +10,8 @@
aapt_version = "aapt",
assets = [
"//tensorflow/lite/java/demo/app/src/main/assets:labels_mobilenet_quant_v1_224.txt",
- "@tflite_mobilenet//:mobilenet_quant_v1_224.tflite",
+ "@tflite_mobilenet_quant//:mobilenet_v1_1.0_224_quant.tflite",
+ "@tflite_mobilenet_float//:mobilenet_v1_1.0_224.tflite",
],
assets_dir = "",
custom_package = "com.example.android.tflitecamerademo",
diff --git a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java
index 3596e42..165d335 100644
--- a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java
+++ b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/Camera2BasicFragment.java
@@ -56,11 +56,12 @@
import android.view.TextureView;
import android.view.View;
import android.view.ViewGroup;
-import android.widget.CompoundButton;
+import android.widget.AdapterView;
+import android.widget.ArrayAdapter;
+import android.widget.ListView;
import android.widget.NumberPicker;
import android.widget.TextView;
import android.widget.Toast;
-import android.widget.ToggleButton;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
@@ -70,6 +71,7 @@
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
+
/** Basic fragments for the Camera. */
public class Camera2BasicFragment extends Fragment
implements FragmentCompat.OnRequestPermissionsResultCallback {
@@ -87,9 +89,11 @@
private boolean runClassifier = false;
private boolean checkedPermissions = false;
private TextView textView;
- private ToggleButton toggle;
private NumberPicker np;
private ImageClassifier classifier;
+ private ListView deviceView;
+ private ListView modelView;
+
/** Max preview width that is guaranteed by Camera2 API */
private static final int MAX_PREVIEW_WIDTH = 1920;
@@ -123,6 +127,15 @@
public void onSurfaceTextureUpdated(SurfaceTexture texture) {}
};
+ // Model parameter constants.
+ private String gpu;
+ private String cpu;
+ private String nnApi;
+ private String mobilenetV1Quant;
+ private String mobilenetV1Float;
+
+
+
/** ID of the current {@link CameraDevice}. */
private String cameraId;
@@ -169,6 +182,14 @@
}
};
+ private ArrayList<String> deviceStrings = new ArrayList<String>();
+ private ArrayList<String> modelStrings = new ArrayList<String>();
+
+ /** Current indices of device and model. */
+ int currentDevice = -1;
+
+ int currentModel = -1;
+
/** An additional thread for running tasks that shouldn't block the UI. */
private HandlerThread backgroundThread;
@@ -298,17 +319,113 @@
return inflater.inflate(R.layout.fragment_camera2_basic, container, false);
}
+ private void updateActiveModel() {
+ // Get UI information before delegating to background
+ final int modelIndex = modelView.getCheckedItemPosition();
+ final int deviceIndex = deviceView.getCheckedItemPosition();
+
+ backgroundHandler.post(() -> {
+ if (modelIndex == currentModel && deviceIndex == currentDevice) {
+ return;
+ }
+ currentModel = modelIndex;
+ currentDevice = deviceIndex;
+
+ // Disable classifier while updating
+ if (classifier != null) {
+ classifier.close();
+ classifier = null;
+ }
+
+ // Lookup names of parameters.
+ String model = modelStrings.get(modelIndex);
+ String device = deviceStrings.get(deviceIndex);
+
+ Log.i(TAG, "Changing model to " + model + " device " + device);
+
+ // Try to load model.
+ try {
+ if (model.equals(mobilenetV1Quant)) {
+ classifier = new ImageClassifierQuantizedMobileNet(getActivity());
+ } else if (model.equals(mobilenetV1Float)) {
+ classifier = new ImageClassifierFloatMobileNet(getActivity());
+ } else {
+ showToast("Failed to load model");
+ }
+ } catch (IOException e) {
+ Log.d(TAG, "Failed to load", e);
+ classifier = null;
+ }
+
+ // Customzie the interpreter to the type of device we want to use.
+ if (device.equals(cpu)) {
+ } else if (device.equals(gpu)) {
+ if (!GpuDelegateHelper.isGpuDelegateAvailable()) {
+ showToast("gpu not in this build.");
+ classifier = null;
+ } else if (model.equals(mobilenetV1Quant)) {
+ showToast("gpu requires float model.");
+ classifier = null;
+ } else {
+ classifier.useGpu();
+ }
+ } else if (device.equals(nnApi)) {
+ classifier.useNNAPI();
+ }
+ });
+ }
+
/** Connect the buttons to their event handler. */
@Override
public void onViewCreated(final View view, Bundle savedInstanceState) {
+ gpu = getString(R.string.gpu);
+ cpu = getString(R.string.cpu);
+ nnApi = getString(R.string.nnapi);
+ mobilenetV1Quant = getString(R.string.mobilenetV1Quant);
+ mobilenetV1Float = getString(R.string.mobilenetV1Float);
+
+ // Get references to widgets.
textureView = (AutoFitTextureView) view.findViewById(R.id.texture);
textView = (TextView) view.findViewById(R.id.text);
- toggle = (ToggleButton) view.findViewById(R.id.button);
+ deviceView = (ListView) view.findViewById(R.id.device);
+ modelView = (ListView) view.findViewById(R.id.model);
- toggle.setOnCheckedChangeListener(
- new CompoundButton.OnCheckedChangeListener() {
- public void onCheckedChanged(CompoundButton buttonView, boolean isChecked) {
- backgroundHandler.post(() -> classifier.setUseNNAPI(isChecked));
+ // Build list of models
+ modelStrings.add(mobilenetV1Quant);
+ modelStrings.add(mobilenetV1Float);
+
+ // Build list of devices
+ int defaultModelIndex = 0;
+ deviceStrings.add(cpu);
+ if (GpuDelegateHelper.isGpuDelegateAvailable()) {
+ deviceStrings.add(gpu);
+ }
+ deviceStrings.add(nnApi);
+
+ deviceView.setAdapter(
+ new ArrayAdapter<String>(
+ getContext(), R.layout.listview_row, R.id.listview_row_text, deviceStrings));
+ deviceView.setChoiceMode(ListView.CHOICE_MODE_SINGLE);
+ deviceView.setOnItemClickListener(
+ new AdapterView.OnItemClickListener() {
+ @Override
+ public void onItemClick(AdapterView<?> parent, View view, int position, long id) {
+ updateActiveModel();
+ }
+ });
+ deviceView.setItemChecked(0, true);
+
+ modelView.setChoiceMode(ListView.CHOICE_MODE_SINGLE);
+ ArrayAdapter<String> modelAdapter =
+ new ArrayAdapter<>(
+ getContext(), R.layout.listview_row, R.id.listview_row_text, modelStrings);
+ modelView.setAdapter(modelAdapter);
+ modelView.setItemChecked(defaultModelIndex, true);
+ modelView.setOnItemClickListener(
+ new AdapterView.OnItemClickListener() {
+ @Override
+ public void onItemClick(AdapterView<?> parent, View view, int position, long id) {
+ updateActiveModel();
}
});
@@ -323,18 +440,14 @@
backgroundHandler.post(() -> classifier.setNumThreads(newVal));
}
});
+
+ // Start initial model.
}
/** Load the model and labels. */
@Override
public void onActivityCreated(Bundle savedInstanceState) {
super.onActivityCreated(savedInstanceState);
- try {
- // create either a new ImageClassifierQuantizedMobileNet or an ImageClassifierFloatInception
- classifier = new ImageClassifierQuantizedMobileNet(getActivity());
- } catch (IOException e) {
- Log.e(TAG, "Failed to initialize an image classifier.", e);
- }
startBackgroundThread();
}
@@ -562,10 +675,12 @@
backgroundThread = new HandlerThread(HANDLE_THREAD_NAME);
backgroundThread.start();
backgroundHandler = new Handler(backgroundThread.getLooper());
+ // Start the classification train & load an initial model.
synchronized (lock) {
runClassifier = true;
}
backgroundHandler.post(periodicClassify);
+ updateActiveModel();
}
/** Stops the background thread and its {@link Handler}. */
diff --git a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/GpuDelegateHelper.java b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/GpuDelegateHelper.java
new file mode 100644
index 0000000..8dca177
--- /dev/null
+++ b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/GpuDelegateHelper.java
@@ -0,0 +1,49 @@
+/* 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.
+==============================================================================*/
+
+package com.example.android.tflitecamerademo;
+
+import org.tensorflow.lite.Delegate;
+
+/**
+ * Helper class for {@code GpuDelegate}.
+ *
+ * <p>WARNING: This is an experimental API and subject to change.
+ */
+public class GpuDelegateHelper {
+ private GpuDelegateHelper() {}
+
+ /** Checks whether {@code GpuDelegate} is available. */
+ public static boolean isGpuDelegateAvailable() {
+ try {
+ Class.forName("org.tensorflow.lite.experimental.GpuDelegate");
+ return true;
+ } catch (Exception e) {
+ return false;
+ }
+ }
+
+ /** Returns an instance of {@code GpuDelegate} if available. */
+ public static Delegate createGpuDelegate() {
+ try {
+ return Class.forName("org.tensorflow.lite.experimental.GpuDelegate")
+ .asSubclass(Delegate.class)
+ .getDeclaredConstructor()
+ .newInstance();
+ } catch (Exception e) {
+ throw new IllegalStateException(e);
+ }
+ }
+}
diff --git a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java
index 39057aa..512f8b6 100644
--- a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java
+++ b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifier.java
@@ -38,6 +38,7 @@
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
+import org.tensorflow.lite.Delegate;
import org.tensorflow.lite.Interpreter;
/**
@@ -93,6 +94,9 @@
}
});
+ /** holds a gpu delegate */
+ Delegate gpuDelegate = null;
+
/** Initializes an {@code ImageClassifier}. */
ImageClassifier(Activity activity) throws IOException {
tfliteModel = loadModelFile(activity);
@@ -159,12 +163,27 @@
private void recreateInterpreter() {
if (tflite != null) {
tflite.close();
+ // TODO(b/120679982)
+ // gpuDelegate.close();
tflite = new Interpreter(tfliteModel, tfliteOptions);
}
}
- public void setUseNNAPI(Boolean nnapi) {
- tfliteOptions.setUseNNAPI(nnapi);
+ public void useGpu() {
+ if (gpuDelegate == null && GpuDelegateHelper.isGpuDelegateAvailable()) {
+ gpuDelegate = GpuDelegateHelper.createGpuDelegate();
+ tfliteOptions.addDelegate(gpuDelegate);
+ recreateInterpreter();
+ }
+ }
+
+ public void useCPU() {
+ tfliteOptions.setUseNNAPI(false);
+ recreateInterpreter();
+ }
+
+ public void useNNAPI() {
+ tfliteOptions.setUseNNAPI(true);
recreateInterpreter();
}
diff --git a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierFloatMobileNet.java b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierFloatMobileNet.java
new file mode 100644
index 0000000..c87ffff
--- /dev/null
+++ b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierFloatMobileNet.java
@@ -0,0 +1,94 @@
+/* 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.
+==============================================================================*/
+
+package com.example.android.tflitecamerademo;
+
+import android.app.Activity;
+import java.io.IOException;
+
+/** This classifier works with the float MobileNet model. */
+public class ImageClassifierFloatMobileNet extends ImageClassifier {
+
+ /**
+ * An array to hold inference results, to be feed into Tensorflow Lite as outputs. This isn't part
+ * of the super class, because we need a primitive array here.
+ */
+ private float[][] labelProbArray = null;
+
+ /**
+ * Initializes an {@code ImageClassifierFloatMobileNet}.
+ *
+ * @param activity
+ */
+ ImageClassifierFloatMobileNet(Activity activity) throws IOException {
+ super(activity);
+ labelProbArray = new float[1][getNumLabels()];
+ }
+
+ @Override
+ protected String getModelPath() {
+ // you can download this file from
+ // see build.gradle for where to obtain this file. It should be auto
+ // downloaded into assets.
+ return "mobilenet_v1_1.0_224.tflite";
+ }
+
+ @Override
+ protected String getLabelPath() {
+ return "labels_mobilenet_quant_v1_224.txt";
+ }
+
+ @Override
+ protected int getImageSizeX() {
+ return 224;
+ }
+
+ @Override
+ protected int getImageSizeY() {
+ return 224;
+ }
+
+ @Override
+ protected int getNumBytesPerChannel() {
+ return 4; // Float.SIZE / Byte.SIZE;
+ }
+
+ @Override
+ protected void addPixelValue(int pixelValue) {
+ imgData.putFloat(((pixelValue >> 16) & 0xFF) / 255.f);
+ imgData.putFloat(((pixelValue >> 8) & 0xFF) / 255.f);
+ imgData.putFloat((pixelValue & 0xFF) / 255.f);
+ }
+
+ @Override
+ protected float getProbability(int labelIndex) {
+ return labelProbArray[0][labelIndex];
+ }
+
+ @Override
+ protected void setProbability(int labelIndex, Number value) {
+ labelProbArray[0][labelIndex] = value.floatValue();
+ }
+
+ @Override
+ protected float getNormalizedProbability(int labelIndex) {
+ return labelProbArray[0][labelIndex];
+ }
+
+ @Override
+ protected void runInference() {
+ tflite.run(imgData, labelProbArray);
+ }
+}
diff --git a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java
index e164ac7..6310a56 100644
--- a/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java
+++ b/tensorflow/lite/java/demo/app/src/main/java/com/example/android/tflitecamerademo/ImageClassifierQuantizedMobileNet.java
@@ -42,8 +42,9 @@
@Override
protected String getModelPath() {
// you can download this file from
- // https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip
- return "mobilenet_quant_v1_224.tflite";
+ // see build.gradle for where to obtain this file. It should be auto
+ // downloaded into assets.
+ return "mobilenet_v1_1.0_224_quant.tflite";
}
@Override
diff --git a/tensorflow/lite/java/demo/app/src/main/res/drawable/item_selector.xml b/tensorflow/lite/java/demo/app/src/main/res/drawable/item_selector.xml
new file mode 100644
index 0000000..202c900
--- /dev/null
+++ b/tensorflow/lite/java/demo/app/src/main/res/drawable/item_selector.xml
@@ -0,0 +1,11 @@
+<?xml version="1.0" encoding="utf-8"?>
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+
+ <!-- pressed -->
+ <item android:drawable="@color/selection_highlight" android:state_pressed="true" />
+ <!-- focused -->
+ <item android:drawable="@color/selection_focus" android:state_activated="true" />
+ <!-- default -->
+ <item android:drawable="@color/item_normal" />
+
+</selector>
diff --git a/tensorflow/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml b/tensorflow/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml
index ef8a9e0..ee71ab8 100644
--- a/tensorflow/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml
+++ b/tensorflow/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml
@@ -40,12 +40,27 @@
android:scaleType="centerInside"
android:src="@drawable/logo"/>
- <ToggleButton
- android:id="@+id/button"
+ <RadioGroup
+ android:gravity="center"
android:layout_width="match_parent"
- android:layout_height="wrap_content"
- android:textOff="@string/tflite"
- android:textOn="@string/nnapi"/>
+ android:layout_height="match_parent"
+ android:orientation="horizontal">
+ <RadioButton
+ android:id="@+id/radio_cpu"
+ android:background="#0000000f"
+ android:layout_width="wrap_content"
+ android:layout_height="wrap_content"
+ android:text="@string/cpu"
+ android:textColor="@android:color/white" />
+ <RadioButton
+ android:id="@+id/radio_nnapi"
+ android:background="#0000000f"
+ android:layout_width="wrap_content"
+ android:layout_height="wrap_content"
+ android:text="@string/nnapi"
+ android:textColor="@android:color/white" />
+ </RadioGroup>
+
<NumberPicker
android:id="@+id/np"
android:layout_width="wrap_content"
diff --git a/tensorflow/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml b/tensorflow/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml
index ddb099a..70eedfd 100644
--- a/tensorflow/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml
+++ b/tensorflow/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml
@@ -57,38 +57,83 @@
android:textStyle="bold" />
</LinearLayout>
- <LinearLayout
- android:orientation="horizontal"
- android:background="#513400"
- android:layout_alignParentBottom="true"
- android:layout_width="match_parent"
+ <LinearLayout
android:id="@+id/bottom_info_view"
+ android:layout_width="match_parent"
+ android:layout_height="200dp"
+
+ android:layout_alignParentBottom="true"
android:layout_marginBottom="10dp"
- android:layout_height="50dp">
- <TextView
+ android:background="#513400"
+ android:orientation="horizontal">
+
+ <LinearLayout
android:layout_width="wrap_content"
android:layout_height="match_parent"
- android:textColor="@android:color/white"
- android:textAlignment="center"
- android:gravity="center"
- android:text="Threads:"/>
- <NumberPicker
- android:id="@+id/np"
- android:layout_width="wrap_content"
- android:layout_height="wrap_content"
- android:layout_marginLeft="10dp"
- android:theme="@style/AppTheme.Picker"
- android:visibility="visible" />
- <ToggleButton
- android:id="@+id/button"
- android:textOff="@string/tflite"
- android:textOn="@string/nnapi"
- android:layout_width="wrap_content"
- android:layout_height="wrap_content"
- android:layout_marginLeft="10dp"
- android:background="#0000000f"
- android:textColor="@android:color/white" />
+ android:orientation="vertical">
+
+ <TextView
+ android:layout_width="wrap_content"
+ android:layout_height="wrap_content"
+ android:gravity="center"
+ android:text="Threads"
+ android:textAlignment="center"
+ android:textColor="@android:color/white" />
+
+ <NumberPicker
+ android:id="@+id/np"
+ android:layout_width="wrap_content"
+ android:layout_height="wrap_content"
+ android:layout_marginLeft="10dp"
+ android:theme="@style/AppTheme.Picker"
+ android:visibility="visible" />
+
+ </LinearLayout>
+
+ <LinearLayout
+ android:id="@+id/modelLayout"
+ android:layout_width="150dp"
+ android:layout_height="match_parent"
+ android:orientation="vertical">
+
+ <TextView
+ android:id="@+id/textView"
+ android:layout_width="match_parent"
+ android:layout_height="20dp"
+ android:text="@string/modelLabel"
+ android:textAlignment="center"
+ android:textColor="@android:color/white" />
+
+ <ListView
+ android:id="@+id/model"
+ android:layout_width="match_parent"
+ android:layout_height="180dp">
+
+ </ListView>
+ </LinearLayout>
+
+ <LinearLayout
+ android:id="@+id/deviceLayout"
+ android:layout_width="140dp"
+ android:layout_height="match_parent"
+ android:orientation="vertical">
+
+ <TextView
+ android:id="@+id/textView2"
+ android:layout_width="match_parent"
+ android:layout_height="20dp"
+ android:text="@string/deviceLabel"
+ android:textAlignment="center"
+ android:textColor="@android:color/white" />
+
+ <ListView
+ android:id="@+id/device"
+ android:layout_width="match_parent"
+ android:layout_height="180dp" />
+
+ </LinearLayout>
+
</LinearLayout>
diff --git a/tensorflow/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml b/tensorflow/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
index e567009..f8312cc 100644
--- a/tensorflow/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
+++ b/tensorflow/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
@@ -57,22 +57,30 @@
android:textStyle="bold" />
</LinearLayout>
- <LinearLayout
- android:orientation="horizontal"
- android:background="#aa7700"
- android:layout_alignParentBottom="true"
- android:layout_width="match_parent"
+ <LinearLayout
android:id="@+id/bottom_info_view"
+ android:layout_width="match_parent"
+ android:layout_height="200dp"
+
+ android:layout_alignParentBottom="true"
android:layout_marginBottom="10dp"
- android:layout_height="50dp">
- <TextView
+ android:background="#513400"
+ android:orientation="horizontal">
+
+ <LinearLayout
android:layout_width="wrap_content"
android:layout_height="match_parent"
- android:textColor="@android:color/white"
- android:textAlignment="center"
+ android:orientation="vertical">
+
+ <TextView
+ android:layout_width="wrap_content"
+ android:layout_height="wrap_content"
android:gravity="center"
- android:text="@string/threads" />
+ android:text="Threads"
+ android:textAlignment="center"
+ android:textColor="@android:color/white" />
+
<NumberPicker
android:id="@+id/np"
android:layout_width="wrap_content"
@@ -80,15 +88,51 @@
android:layout_marginLeft="10dp"
android:theme="@style/AppTheme.Picker"
android:visibility="visible" />
- <ToggleButton
- android:id="@+id/button"
- android:textOff="@string/tflite"
- android:textOn="@string/nnapi"
- android:layout_width="wrap_content"
- android:layout_height="wrap_content"
- android:layout_marginLeft="10dp"
- android:background="#0000000f"
- android:textColor="@android:color/white" />
+
+ </LinearLayout>
+
+ <LinearLayout
+ android:id="@+id/modelLayout"
+ android:layout_width="150dp"
+ android:layout_height="match_parent"
+ android:orientation="vertical">
+
+ <TextView
+ android:id="@+id/textView"
+ android:layout_width="match_parent"
+ android:layout_height="20dp"
+ android:text="@string/modelLabel"
+ android:textAlignment="center"
+ android:textColor="@android:color/white" />
+
+ <ListView
+ android:id="@+id/model"
+ android:layout_width="match_parent"
+ android:layout_height="180dp">
+
+ </ListView>
+ </LinearLayout>
+
+ <LinearLayout
+ android:id="@+id/deviceLayout"
+ android:layout_width="140dp"
+ android:layout_height="match_parent"
+ android:orientation="vertical">
+
+ <TextView
+ android:id="@+id/textView2"
+ android:layout_width="match_parent"
+ android:layout_height="20dp"
+ android:text="@string/deviceLabel"
+ android:textAlignment="center"
+ android:textColor="@android:color/white" />
+
+ <ListView
+ android:id="@+id/device"
+ android:layout_width="match_parent"
+ android:layout_height="180dp" />
+
+ </LinearLayout>
</LinearLayout>
</RelativeLayout>
diff --git a/tensorflow/lite/java/demo/app/src/main/res/layout/listview_row.xml b/tensorflow/lite/java/demo/app/src/main/res/layout/listview_row.xml
new file mode 100644
index 0000000..349b0f6
--- /dev/null
+++ b/tensorflow/lite/java/demo/app/src/main/res/layout/listview_row.xml
@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+ android:layout_width="match_parent"
+ android:layout_height="match_parent">
+
+
+ <TextView
+ android:id="@+id/listview_row_text"
+ android:layout_width="match_parent"
+ android:layout_height="match_parent"
+ android:layout_marginRight="2dp"
+ android:background="@drawable/item_selector"
+ android:padding="10dp"
+ android:textSize="18sp"
+ android:textStyle="bold" />
+
+</LinearLayout>
\ No newline at end of file
diff --git a/tensorflow/lite/java/demo/app/src/main/res/values/colors.xml b/tensorflow/lite/java/demo/app/src/main/res/values/colors.xml
index 4b75d2b..c30f1dc 100644
--- a/tensorflow/lite/java/demo/app/src/main/res/values/colors.xml
+++ b/tensorflow/lite/java/demo/app/src/main/res/values/colors.xml
@@ -16,4 +16,7 @@
-->
<resources>
<color name="control_background">#cc4285f4</color>
+ <color name="selection_highlight">#aaaaaa</color>
+ <color name="selection_focus">#eeaa55</color>
+ <color name="item_normal">#eeeeee</color>
</resources>
diff --git a/tensorflow/lite/java/demo/app/src/main/res/values/strings.xml b/tensorflow/lite/java/demo/app/src/main/res/values/strings.xml
index 29a033b..8cc88f2 100644
--- a/tensorflow/lite/java/demo/app/src/main/res/values/strings.xml
+++ b/tensorflow/lite/java/demo/app/src/main/res/values/strings.xml
@@ -23,4 +23,11 @@
<string name="toggle">Use NNAPI</string>
<string name="tflite">tflite</string>
<string name="nnapi">NNAPI</string>
+ <string name="gpu">GPU</string>
+ <string name="cpu">CPU</string>
+ <string name="modelLabel">Model</string>
+ <string name="deviceLabel">Device</string>
+ <string name="mobilenetV1Quant">mobilenet v1 quant</string>;
+ <string name="mobilenetV1Float">mobilenet v1 float</string>;;
+
</resources>
diff --git a/tensorflow/lite/java/jni/BUILD b/tensorflow/lite/java/jni/BUILD
new file mode 100644
index 0000000..ce17ac4
--- /dev/null
+++ b/tensorflow/lite/java/jni/BUILD
@@ -0,0 +1,47 @@
+package(default_visibility = ["//tensorflow/lite:__subpackages__"])
+
+licenses(["notice"]) # Apache 2.0
+
+# Helper target for exposing JNI headers across multiple platforms.
+cc_library(
+ name = "jni",
+ hdrs = select({
+ # The Android toolchain makes "jni.h" available in the include path.
+ # For non-Android toolchains, generate jni.h and jni_md.h.
+ "//tensorflow:android": [],
+ "//conditions:default": [
+ ":jni.h",
+ ":jni_md.h",
+ ],
+ }),
+ includes = select({
+ "//tensorflow:android": [],
+ "//conditions:default": ["."],
+ }),
+)
+
+# Silly rules to make
+# #include <jni.h>
+# in the source headers work
+# (in combination with the "includes" attribute of the tf_cuda_library rule
+# above. Not needed when using the Android toolchain).
+#
+# Inspired from:
+# https://github.com/bazelbuild/bazel/blob/f99a0543f8d97339d32075c7176b79f35be84606/src/main/native/BUILD
+# but hopefully there is a simpler alternative to this.
+genrule(
+ name = "copy_jni_h",
+ srcs = ["@bazel_tools//tools/jdk:jni_header"],
+ outs = ["jni.h"],
+ cmd = "cp -f $< $@",
+)
+
+genrule(
+ name = "copy_jni_md_h",
+ srcs = select({
+ "//tensorflow:darwin": ["@bazel_tools//tools/jdk:jni_md_header-darwin"],
+ "//conditions:default": ["@bazel_tools//tools/jdk:jni_md_header-linux"],
+ }),
+ outs = ["jni_md.h"],
+ cmd = "cp -f $< $@",
+)
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 a03d7b5..2203d5f 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
@@ -43,15 +43,34 @@
* <pre>{@code
* Object[] inputs = {input0, input1, ...};
* Map<Integer, Object> map_of_indices_to_outputs = new HashMap<>();
- * float[][][] ith_output = new float[3][2][4];
+ * ByteBuffer ith_output = ByteBuffer.allocateDirect(3 * 2 * 4 * 4); // Float tensor, shape 3x2x4.
+ * ith_output.order(ByteOrder.nativeOrder());
* map_of_indices_to_outputs.put(i, ith_output);
* try (Interpreter interpreter = new Interpreter(file_of_a_tensorflowlite_model)) {
* interpreter.runForMultipleInputsOutputs(inputs, map_of_indices_to_outputs);
* }
* }</pre>
*
+ * <p>If a model takes or produces string tensors:
+ *
+ * <pre>{@code
+ * String[] input = {"foo", "bar"}; // Input tensor shape is [2].
+ * String[] output = new String[3][2]; // Output tensor shape is [3, 2].
+ * try (Interpreter interpreter = new Interpreter(file_of_a_tensorflowlite_model)) {
+ * interpreter.runForMultipleInputsOutputs(input, output);
+ * }
+ * }</pre>
+ *
* <p>Orders of inputs and outputs are determined when converting TensorFlow model to TensorFlowLite
- * model with Toco.
+ * model with Toco, as are the default shapes of the inputs.
+ *
+ * <p>When inputs are provided as (multi-dimensional) arrays, the corresponding input tensor(s) will
+ * be implicitly resized according to that array's shape. When inputs are provided as {@link
+ * ByteBuffer} types, no implicit resizing is done; the caller must ensure that the {@link
+ * ByteBuffer} byte size either matches that of the corresponding tensor, or that they first resize
+ * the tensor via {@link #resizeInput()}. Tensor shape and type information can be obtained via the
+ * {@link Tensor} class, available via {@link #getInputTensor(int)} and {@link
+ * #getOutputTensor(int)}.
*
* <p><b>WARNING:</b>Instances of a {@code Interpreter} is <b>not</b> thread-safe. A {@code
* Interpreter} owns resources that <b>must</b> be explicitly freed by invoking {@link #close()}
@@ -192,12 +211,13 @@
* Runs model inference if the model takes only one input, and provides only one output.
*
* <p>Warning: The API runs much faster if {@link ByteBuffer} is used as input data type. Please
- * consider using {@link ByteBuffer} to feed input data for better performance.
+ * consider using {@link ByteBuffer} to feed primitive input data for better performance.
*
* @param input an array or multidimensional array, or a {@link ByteBuffer} of primitive types
* including int, float, long, and byte. {@link ByteBuffer} is the preferred way to pass large
- * input data. When {@link ByteBuffer} is used, its content should remain unchanged until
- * model inference is done.
+ * input data for primitive types, whereas string types require using the (multi-dimensional)
+ * array input path. When {@link ByteBuffer} is used, its content should remain unchanged
+ * until model inference is done.
* @param output a multidimensional array of output data, or a {@link ByteBuffer} of primitive
* types including int, float, long, and byte.
*/
@@ -212,13 +232,14 @@
* Runs model inference if the model takes multiple inputs, or returns multiple outputs.
*
* <p>Warning: The API runs much faster if {@link ByteBuffer} is used as input data type. Please
- * consider using {@link ByteBuffer} to feed input data for better performance.
+ * consider using {@link ByteBuffer} to feed primitive input data for better performance.
*
* @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 ByteBuffer} of
* primitive types including int, float, long, and byte. {@link ByteBuffer} is the preferred
- * way to pass large input data. When {@link ByteBuffer} is used, its content should remain
- * unchanged until model inference is done.
+ * way to pass large input data, whereas string types require using the (multi-dimensional)
+ * array input path. When {@link 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
* ByteBuffer}s of primitive types including int, float, long, and byte. It only needs to keep
* entries for the outputs to be used.
diff --git a/tensorflow/lite/java/src/main/native/BUILD b/tensorflow/lite/java/src/main/native/BUILD
index 8f95f14..52194e8 100644
--- a/tensorflow/lite/java/src/main/native/BUILD
+++ b/tensorflow/lite/java/src/main/native/BUILD
@@ -15,15 +15,7 @@
"nativeinterpreterwrapper_jni.cc",
"tensor_jni.cc",
"tensorflow_lite_jni.cc",
- ] + select({
- # The Android toolchain makes "jni.h" available in the include path.
- # For non-Android toolchains, generate jni.h and jni_md.h.
- "//tensorflow:android": [],
- "//conditions:default": [
- ":jni.h",
- ":jni_md.h",
- ],
- }),
+ ],
hdrs = [
"exception_jni.h",
"nativeinterpreterwrapper_jni.h",
@@ -31,75 +23,31 @@
"tensorflow_lite_jni.h",
],
copts = tflite_copts(),
- includes = select({
- "//tensorflow:android": [],
- "//conditions:default": ["."],
- }),
linkopts = [
"-lm",
"-ldl",
],
deps = [
- "//tensorflow/lite:context",
"//tensorflow/lite:framework",
"//tensorflow/lite:schema_fbs_version",
"//tensorflow/lite:string_util",
+ "//tensorflow/lite/c:c_api_internal",
+ "//tensorflow/lite/java/jni",
],
alwayslink = 1,
)
-# Silly rules to make
-# #include <jni.h>
-# in the source headers work
-# (in combination with the "includes" attribute of the tf_cuda_library rule
-# above. Not needed when using the Android toolchain).
-#
-# Inspired from:
-# https://github.com/bazelbuild/bazel/blob/f99a0543f8d97339d32075c7176b79f35be84606/src/main/native/BUILD
-# but hopefully there is a simpler alternative to this.
-genrule(
- name = "copy_jni_h",
- srcs = ["@bazel_tools//tools/jdk:jni_header"],
- outs = ["jni.h"],
- cmd = "cp -f $< $@",
-)
-
-genrule(
- name = "copy_jni_md_h",
- srcs = select({
- "//tensorflow:darwin": ["@bazel_tools//tools/jdk:jni_md_header-darwin"],
- "//conditions:default": ["@bazel_tools//tools/jdk:jni_md_header-linux"],
- }),
- outs = ["jni_md.h"],
- cmd = "cp -f $< $@",
-)
-
cc_library(
name = "init_tensorflow",
srcs = [
"init_tensorflow_jni.cc",
- ] + select({
- # The Android toolchain makes "jni.h" available in the include path.
- # For non-Android toolchains, generate jni.h and jni_md.h.
- "//tensorflow:android": [],
- "//conditions:default": [
- ":jni.h",
- ":jni_md.h",
- ],
- }),
+ ],
hdrs = [
"init_tensorflow_jni.h",
],
copts = tflite_copts(),
- includes = select({
- "//tensorflow:android": [],
- "//conditions:default": ["."],
- }),
- linkopts = [
- "-lm",
- "-ldl",
- ],
deps = [
+ "//tensorflow/lite/java/jni",
"//tensorflow/lite/testing:init_tensorflow",
],
alwayslink = 1,
diff --git a/tensorflow/lite/java/src/test/native/BUILD b/tensorflow/lite/java/src/test/native/BUILD
index 4d3e82b..481aea7 100644
--- a/tensorflow/lite/java/src/test/native/BUILD
+++ b/tensorflow/lite/java/src/test/native/BUILD
@@ -12,20 +12,11 @@
testonly = 1,
srcs = [
"interpreter_test_jni.cc",
- ] + select({
- # The Android toolchain makes "jni.h" available in the include path.
- # For non-Android toolchains, generate jni.h and jni_md.h.
- "//tensorflow:android": [],
- "//conditions:default": [
- "//tensorflow/lite/java/src/main/native:jni.h",
- "//tensorflow/lite/java/src/main/native:jni_md.h",
- ],
- }),
- includes = select({
- "//tensorflow:android": [],
- "//conditions:default": ["../../main/native/."],
- }),
- deps = ["//tensorflow/lite/c:c_api_internal"],
+ ],
+ deps = [
+ "//tensorflow/lite/c:c_api_internal",
+ "//tensorflow/lite/java/jni",
+ ],
)
tflite_jni_binary(
diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD
index 00d9d1f..bad1c4a 100644
--- a/tensorflow/lite/kernels/BUILD
+++ b/tensorflow/lite/kernels/BUILD
@@ -26,7 +26,6 @@
size = "small",
srcs = ["optional_tensor_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -124,7 +123,6 @@
size = "small",
srcs = ["kernel_util_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -139,7 +137,6 @@
size = "small",
srcs = ["test_util_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios", # TODO(b/117786830)
],
deps = [
@@ -293,7 +290,6 @@
size = "small",
srcs = ["audio_spectrogram_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -310,7 +306,6 @@
size = "small",
srcs = ["mfcc_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -327,7 +322,6 @@
size = "small",
srcs = ["detection_postprocess_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -344,7 +338,6 @@
size = "small",
srcs = ["relu1_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -361,7 +354,6 @@
size = "small",
srcs = ["sparse_output_fully_connected_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -404,7 +396,6 @@
size = "small",
srcs = ["arg_min_max_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -420,7 +411,6 @@
size = "small",
srcs = ["div_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -436,7 +426,6 @@
size = "small",
srcs = ["sub_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -452,7 +441,6 @@
size = "small",
srcs = ["transpose_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -470,7 +458,6 @@
size = "small",
srcs = ["space_to_batch_nd_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -486,7 +473,6 @@
size = "small",
srcs = ["batch_to_space_nd_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -502,7 +488,6 @@
size = "small",
srcs = ["cast_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -559,7 +544,6 @@
size = "small",
srcs = ["dequantize_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -589,7 +573,6 @@
size = "small",
srcs = ["bidirectional_sequence_lstm_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -606,7 +589,6 @@
size = "small",
srcs = ["floor_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -622,7 +604,6 @@
size = "small",
srcs = ["elementwise_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -638,7 +619,6 @@
size = "small",
srcs = ["unidirectional_sequence_lstm_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -654,7 +634,6 @@
size = "small",
srcs = ["bidirectional_sequence_rnn_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable",
],
deps = [
@@ -670,7 +649,6 @@
size = "small",
srcs = ["unidirectional_sequence_rnn_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -699,7 +677,6 @@
size = "small",
srcs = ["exp_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -715,7 +692,6 @@
size = "small",
srcs = ["fake_quant_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -731,7 +707,6 @@
size = "small",
srcs = ["maximum_minimum_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -747,7 +722,6 @@
size = "small",
srcs = ["reduce_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -776,7 +750,6 @@
size = "small",
srcs = ["pad_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -805,7 +778,6 @@
size = "small",
srcs = ["gather_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -822,7 +794,6 @@
size = "small",
srcs = ["topk_v2_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -959,7 +930,6 @@
size = "small",
srcs = ["log_softmax_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1057,7 +1027,6 @@
size = "small",
srcs = ["split_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1073,7 +1042,6 @@
size = "small",
srcs = ["split_v_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1089,7 +1057,6 @@
size = "small",
srcs = ["squeeze_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1105,7 +1072,6 @@
size = "small",
srcs = ["strided_slice_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1121,7 +1087,6 @@
size = "small",
srcs = ["tile_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1140,7 +1105,6 @@
"comparisons_test.cc",
],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1156,7 +1120,6 @@
size = "small",
srcs = ["neg_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1174,7 +1137,6 @@
"select_test.cc",
],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1192,7 +1154,6 @@
"slice_test.cc",
],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1208,7 +1169,6 @@
size = "small",
srcs = ["transpose_conv_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1225,7 +1185,6 @@
size = "small",
srcs = ["expand_dims_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1242,7 +1201,6 @@
size = "small",
srcs = ["sparse_to_dense_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1259,7 +1217,6 @@
size = "small",
srcs = ["shape_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -1276,7 +1233,6 @@
size = "small",
srcs = ["pow_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index a766542..ab09cf7 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -373,11 +373,8 @@
TfLiteTensor* output = GetOutput(context, node, 0);
switch (input->type) {
case kTfLiteFloat32: {
- size_t elements = input->bytes / sizeof(float);
- float* in = input->data.f;
- float* in_end = in + elements;
- float* out = output->data.f;
- for (; in < in_end; in++, out++) *out = std::tanh(*in);
+ optimized_ops::Tanh(GetTensorShape(input), GetTensorData<float>(input),
+ GetTensorShape(output), GetTensorData<float>(output));
return kTfLiteOk;
} break;
case kTfLiteInt16: {
diff --git a/tensorflow/lite/kernels/dequantize.cc b/tensorflow/lite/kernels/dequantize.cc
index b2825bb..7f03c73 100644
--- a/tensorflow/lite/kernels/dequantize.cc
+++ b/tensorflow/lite/kernels/dequantize.cc
@@ -18,6 +18,7 @@
#include "tensorflow/lite/c/builtin_op_data.h"
#include "tensorflow/lite/c/c_api_internal.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/tensor.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/op_macros.h"
@@ -57,7 +58,8 @@
OpContext op_context(context, node);
- TF_LITE_ENSURE(context, op_context.input->type == kTfLiteUInt8);
+ TF_LITE_ENSURE(context, op_context.input->type == kTfLiteUInt8 ||
+ op_context.input->type == kTfLiteInt8);
op_context.output->type = kTfLiteFloat32;
// If the input tensor is constant, we can persist the dequantized value in
@@ -80,10 +82,25 @@
tflite::DequantizationParams op_params;
op_params.zero_point = op_context.input->params.zero_point;
op_params.scale = op_context.input->params.scale;
- optimized_ops::Dequantize(op_params, GetTensorShape(op_context.input),
- GetTensorData<uint8_t>(op_context.input),
- GetTensorShape(op_context.output),
- GetTensorData<float>(op_context.output));
+ switch (op_context.input->type) {
+ case kTfLiteUInt8:
+ optimized_ops::Dequantize(op_params, GetTensorShape(op_context.input),
+ GetTensorData<uint8_t>(op_context.input),
+ GetTensorShape(op_context.output),
+ GetTensorData<float>(op_context.output));
+ break;
+ case kTfLiteInt8:
+ reference_integer_ops::Dequantize(
+ op_params, GetTensorShape(op_context.input),
+ GetTensorData<int8_t>(op_context.input),
+ GetTensorShape(op_context.output),
+ GetTensorData<float>(op_context.output));
+ break;
+ default:
+ context->ReportError(context, "Type %d not supported.",
+ op_context.input->type);
+ return kTfLiteError;
+ }
if (IsConstantTensor(op_context.input)) {
op_data->float_dequantized_weights_initialized = true;
diff --git a/tensorflow/lite/kernels/dequantize_test.cc b/tensorflow/lite/kernels/dequantize_test.cc
index 55265d9..bb5f1e7 100644
--- a/tensorflow/lite/kernels/dequantize_test.cc
+++ b/tensorflow/lite/kernels/dequantize_test.cc
@@ -25,8 +25,16 @@
class DequantizeOpModel : public SingleOpModel {
public:
- DequantizeOpModel(std::initializer_list<int> shape, float min, float max) {
- input_ = AddInput({TensorType_UINT8, shape, min, max});
+ DequantizeOpModel(TensorType type, std::initializer_list<int> shape,
+ float scale, int32_t zero_point) {
+ TensorData input_tensor_data;
+ input_tensor_data.type = type;
+ input_tensor_data.shape = shape;
+ input_tensor_data.min = 0;
+ input_tensor_data.max = 0;
+ input_tensor_data.scale = scale;
+ input_tensor_data.zero_point = zero_point;
+ input_ = AddInput(input_tensor_data);
output_ = AddOutput({TensorType_FLOAT32, shape});
SetBuiltinOp(BuiltinOperator_DEQUANTIZE, BuiltinOptions_DequantizeOptions,
CreateDequantizeOptions(builder_).Union());
@@ -34,7 +42,8 @@
BuildInterpreter({GetShape(input_)});
}
- void SetInput(std::initializer_list<uint8_t> data) {
+ template <typename T>
+ void SetInput(std::initializer_list<T> data) {
PopulateTensor(input_, data);
}
@@ -45,10 +54,22 @@
int output_;
};
-TEST(SplitOpTest, FourDimensional) {
- DequantizeOpModel m({2, 5}, -63.5, 64);
+TEST(DequantizeOpTest, UINT8) {
+ // [-63.5, 64] -> scale=0.5 zero_point=127 for UINT8
+ DequantizeOpModel m(TensorType_UINT8, {2, 5}, 0.5, 127);
- m.SetInput({0, 1, 2, 3, 4, 251, 252, 253, 254, 255});
+ m.SetInput<uint8>({0, 1, 2, 3, 4, 251, 252, 253, 254, 255});
+ m.Invoke();
+ EXPECT_THAT(m.GetOutput(),
+ ElementsAreArray(ArrayFloatNear(
+ {-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64})));
+}
+
+TEST(DequantizeOpTest, INT8) {
+ // [-63.5, 64] -> scale=0.5, zero_point=1 for INT8
+ DequantizeOpModel m(TensorType_INT8, {2, 5}, 0.5, -1);
+
+ m.SetInput<int8>({-128, -127, -126, -125, -124, 123, 124, 125, 126, 127});
m.Invoke();
EXPECT_THAT(m.GetOutput(),
ElementsAreArray(ArrayFloatNear(
diff --git a/tensorflow/lite/kernels/internal/BUILD b/tensorflow/lite/kernels/internal/BUILD
index 7d2653f..6981658 100644
--- a/tensorflow/lite/kernels/internal/BUILD
+++ b/tensorflow/lite/kernels/internal/BUILD
@@ -254,7 +254,6 @@
name = "tensor_test",
srcs = ["tensor_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios", # TODO(b/117786830)
],
deps = [
@@ -287,7 +286,6 @@
name = "quantization_util_test",
srcs = ["quantization_util_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable_ios", # TODO(b/117786830)
],
deps = [
@@ -315,6 +313,7 @@
"reference/depthwiseconv_float.h",
"reference/depthwiseconv_uint8.h",
"reference/fully_connected.h",
+ "reference/integer_ops/dequantize.h",
"reference/reference_ops.h",
"reference/softmax.h",
],
@@ -562,7 +561,6 @@
}),
linkstatic = 1,
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -576,7 +574,6 @@
cc_test(
name = "depthwiseconv_float_test",
srcs = ["depthwiseconv_float_test.cc"],
- tags = ["no_oss"],
deps = [
":optimized_base",
":reference_base",
@@ -591,7 +588,6 @@
srcs = ["depthwiseconv_quantized_test.cc"],
shard_count = 2,
tags = [
- "no_oss",
"tflite_not_portable_ios",
],
deps = [
@@ -607,7 +603,6 @@
name = "resize_bilinear_test",
srcs = ["resize_bilinear_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable",
],
deps = [
@@ -623,7 +618,6 @@
name = "resize_nearest_neighbor_test",
srcs = ["resize_nearest_neighbor_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable",
],
deps = [
@@ -641,7 +635,6 @@
srcs = [
"softmax_quantized_test.cc",
],
- tags = ["no_oss"],
deps = [
":optimized_base",
":quantization_util",
@@ -659,7 +652,6 @@
"logsoftmax_quantized_test.cc",
],
tags = [
- "no_oss",
"tflite_not_portable",
],
deps = [
@@ -675,7 +667,6 @@
cc_test(
name = "log_quantized_test",
srcs = ["log_quantized_test.cc"],
- tags = ["no_oss"],
deps = [
":optimized_base",
":reference_base",
@@ -703,7 +694,6 @@
cc_test(
name = "batch_to_space_nd_test",
srcs = ["batch_to_space_nd_test.cc"],
- tags = ["no_oss"],
deps = [
":optimized_base",
"@com_google_googletest//:gtest_main",
diff --git a/tensorflow/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/lite/kernels/internal/optimized/optimized_ops.h
index c7691e2..c79b69a 100644
--- a/tensorflow/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/lite/kernels/internal/optimized/optimized_ops.h
@@ -3582,8 +3582,8 @@
std::min(params.filter_height, input_height - in_y_origin);
const int filter_count =
(filter_x_end - filter_x_start) * (filter_y_end - filter_y_start);
- // 1280 required by Inception v3
- static constexpr int kAccBufferMaxSize = 2048;
+ // 2560 is required by MobileNetV2 with depth multiplier 2.
+ static constexpr int kAccBufferMaxSize = 4096;
TFLITE_DCHECK_LE(depth, kAccBufferMaxSize);
uint16 acc[kAccBufferMaxSize];
memset(acc, 0, depth * sizeof(acc[0]));
@@ -3748,8 +3748,8 @@
const int filter_y_start = std::max(0, -in_y_origin);
const int filter_y_end =
std::min(params.filter_height, input_height - in_y_origin);
- // 2048 required by Inception v3
- static constexpr int kAccBufferMaxSize = 2048;
+ // 2560 is required by MobileNetV2 with depth multiplier 2.
+ static constexpr int kAccBufferMaxSize = 4096;
TFLITE_DCHECK_LE(depth, kAccBufferMaxSize);
uint8 acc[kAccBufferMaxSize];
memset(acc, 0, depth * sizeof(acc[0]));
diff --git a/tensorflow/lite/kernels/internal/reference/integer_ops/README.md b/tensorflow/lite/kernels/internal/reference/integer_ops/README.md
new file mode 100644
index 0000000..4b1d3c9
--- /dev/null
+++ b/tensorflow/lite/kernels/internal/reference/integer_ops/README.md
@@ -0,0 +1,8 @@
+This directory contains reference implementations for int8 fully integer kernels.
+
+Weight filters of convs are expected to be symmetric per-channel quantized in
+the range [-127, 127].
+Inputs/activations are expected to be asymmetric per-layer quantized in the
+range [-128, 127].
+
+THESE ARE EXPERIMENTAL AND PRONE TO CHANGE.
diff --git a/tensorflow/lite/kernels/internal/reference/integer_ops/dequantize.h b/tensorflow/lite/kernels/internal/reference/integer_ops/dequantize.h
new file mode 100644
index 0000000..03dcb6c
--- /dev/null
+++ b/tensorflow/lite/kernels/internal/reference/integer_ops/dequantize.h
@@ -0,0 +1,42 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_DEQUANTIZE_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_DEQUANTIZE_H_
+
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace reference_integer_ops {
+
+inline void Dequantize(const tflite::DequantizationParams& op_params,
+ const RuntimeShape& input_shape, const int8* input_data,
+ const RuntimeShape& output_shape, float* output_data) {
+ const int32 zero_point = op_params.zero_point;
+ const double scale = op_params.scale;
+ const int flat_size = MatchingFlatSize(input_shape, output_shape);
+
+ for (int i = 0; i < flat_size; i++) {
+ const int32 val = input_data[i];
+ const float result = static_cast<float>(scale * (val - zero_point));
+ output_data[i] = result;
+ }
+}
+
+} // namespace reference_integer_ops
+} // namespace tflite
+
+#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_DEQUANTIZE_H_
diff --git a/tensorflow/lite/kernels/layer_norm_lstm.cc b/tensorflow/lite/kernels/layer_norm_lstm.cc
index 5b0046a..49e8a53 100644
--- a/tensorflow/lite/kernels/layer_norm_lstm.cc
+++ b/tensorflow/lite/kernels/layer_norm_lstm.cc
@@ -55,7 +55,7 @@
constexpr int kCellToOutputWeightsTensor = 11; // Optional
// Layer norm weights tensors of size {n_cell}, representing a diagonal matrix.
-constexpr int kInputLayerNormWeightsTensor = 12;
+constexpr int kInputLayerNormWeightsTensor = 12; // Optional
constexpr int kForgetLayerNormWeightsTensor = 13;
constexpr int kCellLayerNormWeightsTensor = 14;
constexpr int kOutputLayerNormWeightsTensor = 15;
@@ -118,7 +118,8 @@
const TfLiteTensor* input_to_input_weights =
GetOptionalInputTensor(context, node, kInputToInputWeightsTensor);
- if (input_to_input_weights != nullptr) {
+ const bool use_cifg = (input_to_input_weights == nullptr);
+ if (!use_cifg) {
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);
@@ -138,7 +139,9 @@
const TfLiteTensor* recurrent_to_input_weights =
GetOptionalInputTensor(context, node, kRecurrentToInputWeightsTensor);
- if (recurrent_to_input_weights != nullptr) {
+ if (use_cifg) {
+ TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights, nullptr);
+ } else {
TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->size, 2);
TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->data[0],
n_cell);
@@ -161,15 +164,6 @@
TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[1],
n_output);
- // We make sure the input-gate's parameters are either both present (regular
- // LSTM) or not at all (CIFG-LSTM).
- const bool cifg_weights_all_or_none =
- ((input_to_input_weights != nullptr) &&
- (recurrent_to_input_weights != nullptr)) ||
- ((input_to_input_weights == nullptr) &&
- (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);
if (cell_to_input_weights) {
@@ -192,7 +186,6 @@
}
// Making sure the peephole weights are there all or none.
- const bool use_cifg = (input_to_input_weights == nullptr);
const bool peephole_weights_all_or_none =
((cell_to_input_weights != nullptr || use_cifg) &&
(cell_to_forget_weights != nullptr) &&
@@ -204,10 +197,14 @@
// Making sure layer norm weights are not null and have the right dimension.
const TfLiteTensor* input_layer_norm_weights =
- GetInput(context, node, kInputLayerNormWeightsTensor);
- TF_LITE_ENSURE(context, input_layer_norm_weights != nullptr);
- TF_LITE_ENSURE_EQ(context, input_layer_norm_weights->dims->size, 1);
- TF_LITE_ENSURE_EQ(context, input_layer_norm_weights->dims->data[0], n_cell);
+ GetOptionalInputTensor(context, node, kInputLayerNormWeightsTensor);
+ if (use_cifg) {
+ TF_LITE_ENSURE_EQ(context, input_layer_norm_weights, nullptr);
+ } else {
+ TF_LITE_ENSURE(context, input_layer_norm_weights != nullptr);
+ TF_LITE_ENSURE_EQ(context, input_layer_norm_weights->dims->size, 1);
+ TF_LITE_ENSURE_EQ(context, input_layer_norm_weights->dims->data[0], n_cell);
+ }
const TfLiteTensor* forget_layer_norm_weights =
GetInput(context, node, kForgetLayerNormWeightsTensor);
@@ -978,6 +975,9 @@
(projection_weights == nullptr) ? nullptr : projection_weights->data.f;
const float* projection_bias_ptr =
(projection_bias == nullptr) ? nullptr : projection_bias->data.f;
+ const float* input_layer_norm_weight_ptr =
+ (input_layer_norm_weights == nullptr) ? nullptr
+ : input_layer_norm_weights->data.f;
// Required tensors, pointers are non-null.
const float* input_ptr_batch = input->data.f;
@@ -990,7 +990,6 @@
recurrent_to_cell_weights->data.f;
const float* recurrent_to_output_weights_ptr =
recurrent_to_output_weights->data.f;
- const float* input_layer_norm_weight_ptr = input_layer_norm_weights->data.f;
const float* forget_layer_norm_weight_ptr = forget_layer_norm_weights->data.f;
const float* cell_layer_norm_weight_ptr = cell_layer_norm_weights->data.f;
const float* output_layer_norm_weight_ptr = output_layer_norm_weights->data.f;
@@ -1115,6 +1114,9 @@
(projection_weights == nullptr) ? 1.0f : projection_weights->params.scale;
const float* projection_bias_ptr =
(projection_bias == nullptr) ? nullptr : projection_bias->data.f;
+ const float* input_layer_norm_weight_ptr =
+ (input_layer_norm_weights == nullptr) ? nullptr
+ : input_layer_norm_weights->data.f;
// Required tensors, pointers are non-null.
const float* input_ptr_batch = input->data.f;
@@ -1141,7 +1143,6 @@
reinterpret_cast<int8_t*>(recurrent_to_output_weights->data.uint8);
const float recurrent_to_output_weights_scale =
recurrent_to_output_weights->params.scale;
- const float* input_layer_norm_weight_ptr = input_layer_norm_weights->data.f;
const float* forget_layer_norm_weight_ptr = forget_layer_norm_weights->data.f;
const float* cell_layer_norm_weight_ptr = cell_layer_norm_weights->data.f;
const float* output_layer_norm_weight_ptr = output_layer_norm_weights->data.f;
@@ -1221,7 +1222,7 @@
GetOptionalInputTensor(context, node, kCellToOutputWeightsTensor);
const TfLiteTensor* input_layer_norm_weights =
- GetInput(context, node, kInputLayerNormWeightsTensor);
+ GetOptionalInputTensor(context, node, kInputLayerNormWeightsTensor);
const TfLiteTensor* forget_layer_norm_weights =
GetInput(context, node, kForgetLayerNormWeightsTensor);
const TfLiteTensor* cell_layer_norm_weights =
diff --git a/tensorflow/lite/kernels/layer_norm_lstm_test.cc b/tensorflow/lite/kernels/layer_norm_lstm_test.cc
index e89bce5..1c13cee 100644
--- a/tensorflow/lite/kernels/layer_norm_lstm_test.cc
+++ b/tensorflow/lite/kernels/layer_norm_lstm_test.cc
@@ -83,7 +83,11 @@
cell_to_output_weights_ = AddNullInput();
}
- input_layer_norm_weights_ = AddInput(TensorType_FLOAT32);
+ if (use_cifg) {
+ input_layer_norm_weights_ = AddNullInput();
+ } else {
+ input_layer_norm_weights_ = AddInput(TensorType_FLOAT32);
+ }
forget_layer_norm_weights_ = AddInput(TensorType_FLOAT32);
cell_layer_norm_weights_ = AddInput(TensorType_FLOAT32);
output_layer_norm_weights_ = AddInput(TensorType_FLOAT32);
@@ -650,6 +654,223 @@
&layer_norm_lstm);
}
+class CifgPeepholeProjectionNoClippingLayerNormLstmTest
+ : public BaseLayerNormLstmTest {
+ void SetUp() override {
+ input_to_forget_weights_ = {-0.6, -0.1, 0.3, 0.2, 0.9, -0.5, -0.2,
+ -0.4, 0.3, -0.8, -0.4, 0.3, -0.5, -0.4,
+ -0.6, 0.3, -0.4, -0.6, -0.5, -0.5};
+ input_to_cell_weights_ = {-0.4, -0.3, -0.2, -0.1, -0.5, 0.5, -0.2,
+ -0.3, -0.2, -0.6, 0.6, -0.1, -0.4, -0.3,
+ -0.7, 0.7, -0.9, -0.5, 0.8, 0.6};
+ input_to_output_weights_ = {-0.8, -0.4, -0.2, -0.9, -0.1, -0.7, 0.3,
+ -0.3, -0.8, -0.2, 0.6, -0.2, 0.4, -0.7,
+ -0.3, -0.5, 0.1, 0.5, -0.6, -0.4};
+
+ forget_gate_bias_ = {0.1, -0.3, -0.2, 0.1};
+ cell_gate_bias_ = {-0.05, 0.72, 0.25, 0.08};
+ output_gate_bias_ = {0.05, -0.01, 0.2, 0.1};
+
+ recurrent_to_cell_weights_ = {-0.3, 0.2, 0.1, -0.3, 0.8, -0.08,
+ -0.2, 0.3, 0.8, -0.6, -0.1, 0.2};
+ recurrent_to_forget_weights_ = {-0.5, -0.3, -0.5, -0.2, 0.6, 0.4,
+ 0.9, 0.3, -0.1, 0.2, 0.5, 0.2};
+ recurrent_to_output_weights_ = {0.3, -0.1, 0.1, -0.2, -0.5, -0.7,
+ -0.2, -0.6, -0.1, -0.4, -0.7, -0.2};
+
+ cell_to_forget_weights_ = {-0.02, -0.15, -0.25, -0.03};
+ cell_to_output_weights_ = {0.1, -0.1, -0.5, 0.05};
+
+ forget_layer_norm_weights_ = {0.2, 0.2, 0.4, 0.3};
+ cell_layer_norm_weights_ = {0.7, 0.2, 0.3, 0.8};
+ output_layer_norm_weights_ = {0.6, 0.2, 0.2, 0.5};
+ projection_weights_ = {-0.1, 0.2, 0.01, -0.2, 0.1, 0.5,
+ 0.3, 0.08, 0.07, 0.2, -0.4, 0.2};
+
+ layer_norm_lstm_input_ = {
+ {// Batch0: 3 (input_sequence_size) * 5 (n_input)
+ 0.7, 0.8, 0.1, 0.2, 0.3, // seq 0
+ 0.8, 0.1, 0.2, 0.4, 0.5, // seq 1
+ 0.2, 0.7, 0.7, 0.1, 0.7}, // seq 2
+
+ {// Batch1: 3 (input_sequence_size) * 5 (n_input)
+ 0.3, 0.2, 0.9, 0.8, 0.1, // seq 0
+ 0.1, 0.5, 0.2, 0.4, 0.2, // seq 1
+ 0.6, 0.9, 0.2, 0.5, 0.7}, // seq 2
+ };
+ }
+};
+
+TEST_F(CifgPeepholeProjectionNoClippingLayerNormLstmTest,
+ LayerNormLstmBlackBoxTest) {
+ const int n_batch = 2;
+ const int n_input = 5;
+ const int n_cell = 4;
+ const int n_output = 3;
+ const float ceil_clip = 0.0;
+ const float proj_clip = 0.0;
+
+ LayerNormLSTMOpModel layer_norm_lstm(
+ n_batch, n_input, n_cell, n_output,
+ /*use_cifg=*/true, /*use_peephole=*/true,
+ /*use_projection_weights=*/true,
+ /*use_projection_bias=*/false, ceil_clip, proj_clip,
+ {
+ {n_batch, n_input}, // input tensor
+
+ {0, 0}, // input_to_input_weight tensor
+ {n_cell, n_input}, // input_to_forget_weight tensor
+ {n_cell, n_input}, // input_to_cell_weight tensor
+ {n_cell, n_input}, // input_to_output_weight tensor
+
+ {0, 0}, // recurrent_to_input_weight tensor
+ {n_cell, n_output}, // recurrent_to_forget_weight tensor
+ {n_cell, n_output}, // recurrent_to_cell_weight tensor
+ {n_cell, n_output}, // recurrent_to_output_weight tensor
+
+ {0}, // cell_to_input_weight tensor
+ {n_cell}, // cell_to_forget_weight tensor
+ {n_cell}, // cell_to_output_weight tensor
+
+ {0}, // input_layer_norm_weight tensor
+ {n_cell}, // forget_layer_norm_weight tensor
+ {n_cell}, // cell_layer_norm_weight tensor
+ {n_cell}, // output_layer_norm_weight tensor
+
+ {0}, // input_gate_bias tensor
+ {n_cell}, // forget_gate_bias tensor
+ {n_cell}, // cell_bias tensor
+ {n_cell}, // output_gate_bias tensor
+
+ {n_output, n_cell}, // projection_weight tensor
+ {0}, // projection_bias tensor
+ });
+
+ layer_norm_lstm.SetInputToCellWeights(input_to_cell_weights_);
+ layer_norm_lstm.SetInputToForgetWeights(input_to_forget_weights_);
+ layer_norm_lstm.SetInputToOutputWeights(input_to_output_weights_);
+
+ layer_norm_lstm.SetCellBias(cell_gate_bias_);
+ layer_norm_lstm.SetForgetGateBias(forget_gate_bias_);
+ layer_norm_lstm.SetOutputGateBias(output_gate_bias_);
+
+ layer_norm_lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_);
+ layer_norm_lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_);
+ layer_norm_lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_);
+
+ layer_norm_lstm.SetCellToForgetWeights(cell_to_forget_weights_);
+ layer_norm_lstm.SetCellToOutputWeights(cell_to_output_weights_);
+
+ layer_norm_lstm.SetForgetLayerNormWeights(forget_layer_norm_weights_);
+ layer_norm_lstm.SetCellLayerNormWeights(cell_layer_norm_weights_);
+ layer_norm_lstm.SetOutputLayerNormWeights(output_layer_norm_weights_);
+
+ layer_norm_lstm.SetProjectionWeights(projection_weights_);
+
+ // Verify the final output.
+ const std::vector<std::vector<float>> layer_norm_lstm_golden_output = {
+ {
+ // Batch0: 3 (input_sequence_size) * 3 (n_output)
+ 0.02129706, 0.140816242, 0.0112733059, // seq 0
+ 0.0132302344, 0.152308047, 0.0346313119, // seq 1
+ -0.0123688057, 0.165790111, 0.0893077999, // seq 2
+ },
+ {
+ // Batch1: 3 (input_sequence_size) * 3 (n_output)
+ -0.0226350538, 0.0916948169, 0.0769175813, // seq 0
+ -0.0269966982, 0.149707705, 0.094149217, // seq 1
+ -0.0103429332, 0.173016444, 0.0720508844, // seq 2
+ }};
+
+ VerifyGoldens(layer_norm_lstm_input_, layer_norm_lstm_golden_output,
+ &layer_norm_lstm);
+}
+
+TEST_F(CifgPeepholeProjectionNoClippingLayerNormLstmTest,
+ HybridLayerNormLstmBlackBoxTest) {
+ const int n_batch = 2;
+ const int n_input = 5;
+ const int n_cell = 4;
+ const int n_output = 3;
+ const float ceil_clip = 0.0;
+ const float proj_clip = 0.0;
+
+ HybridLayerNormLSTMOpModel layer_norm_lstm(
+ n_batch, n_input, n_cell, n_output,
+ /*use_cifg=*/true, /*use_peephole=*/true,
+ /*use_projection_weights=*/true,
+ /*use_projection_bias=*/false, ceil_clip, proj_clip,
+ {
+ {n_batch, n_input}, // input tensor
+
+ {0, 0}, // input_to_input_weight tensor
+ {n_cell, n_input}, // input_to_forget_weight tensor
+ {n_cell, n_input}, // input_to_cell_weight tensor
+ {n_cell, n_input}, // input_to_output_weight tensor
+
+ {0, 0}, // recurrent_to_input_weight tensor
+ {n_cell, n_output}, // recurrent_to_forget_weight tensor
+ {n_cell, n_output}, // recurrent_to_cell_weight tensor
+ {n_cell, n_output}, // recurrent_to_output_weight tensor
+
+ {0}, // cell_to_input_weight tensor
+ {n_cell}, // cell_to_forget_weight tensor
+ {n_cell}, // cell_to_output_weight tensor
+
+ {0}, // input_layer_norm_weight tensor
+ {n_cell}, // forget_layer_norm_weight tensor
+ {n_cell}, // cell_layer_norm_weight tensor
+ {n_cell}, // output_layer_norm_weight tensor
+
+ {0}, // input_gate_bias tensor
+ {n_cell}, // forget_gate_bias tensor
+ {n_cell}, // cell_bias tensor
+ {n_cell}, // output_gate_bias tensor
+
+ {n_output, n_cell}, // projection_weight tensor
+ {0}, // projection_bias tensor
+ });
+
+ layer_norm_lstm.SetInputToCellWeights(input_to_cell_weights_);
+ layer_norm_lstm.SetInputToForgetWeights(input_to_forget_weights_);
+ layer_norm_lstm.SetInputToOutputWeights(input_to_output_weights_);
+
+ layer_norm_lstm.SetCellBias(cell_gate_bias_);
+ layer_norm_lstm.SetForgetGateBias(forget_gate_bias_);
+ layer_norm_lstm.SetOutputGateBias(output_gate_bias_);
+
+ layer_norm_lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_);
+ layer_norm_lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_);
+ layer_norm_lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_);
+
+ layer_norm_lstm.SetCellToForgetWeights(cell_to_forget_weights_);
+ layer_norm_lstm.SetCellToOutputWeights(cell_to_output_weights_);
+
+ layer_norm_lstm.SetForgetLayerNormWeights(forget_layer_norm_weights_);
+ layer_norm_lstm.SetCellLayerNormWeights(cell_layer_norm_weights_);
+ layer_norm_lstm.SetOutputLayerNormWeights(output_layer_norm_weights_);
+
+ layer_norm_lstm.SetProjectionWeights(projection_weights_);
+
+ // Verify the final output.
+ const std::vector<std::vector<float>> layer_norm_lstm_golden_output = {
+ {
+ // Batch0: 3 (input_sequence_size) * 3 (n_output)
+ 0.0212250091, 0.140474007, 0.0115012666, // seq 0
+ 0.0130806509, 0.152660668, 0.0347516984, // seq 1
+ -0.0124010444, 0.166042402, 0.0898982584, // seq 2
+ },
+ {
+ // Batch1: 3 (input_sequence_size) * 3 (n_output)
+ -0.0228835996, 0.0917588323, 0.0778886303, // seq 0
+ -0.0275101066, 0.148769245, 0.0938384682, // seq 1
+ -0.0103605557, 0.172605693, 0.0728750974, // seq 2
+ }};
+
+ VerifyGoldens(layer_norm_lstm_input_, layer_norm_lstm_golden_output,
+ &layer_norm_lstm);
+}
+
} // namespace
} // namespace custom
} // namespace ops
diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index f179ecb..0c6a462 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -1118,7 +1118,7 @@
cell_to_output_weights_scale, input_gate_bias_ptr,
forget_gate_bias_ptr, cell_bias_ptr, output_gate_bias_ptr,
projection_weights_ptr, projection_weights_scale,
- projection_bias_ptr, params, n_batch, n_cell, n_input,
+ projection_bias_ptr, params, /*n_batch=*/1, n_cell, n_input,
aux_input_size, n_output, output_batch_leading_dim,
input_gate_scratch, forget_gate_scratch, cell_scratch,
output_gate_scratch, scaling_factors_ptr, prod_scaling_factors_ptr,
diff --git a/tensorflow/lite/kernels/pack.cc b/tensorflow/lite/kernels/pack.cc
index 479495c..d15a5a0 100644
--- a/tensorflow/lite/kernels/pack.cc
+++ b/tensorflow/lite/kernels/pack.cc
@@ -35,7 +35,6 @@
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input0 = GetInput(context, node, 0);
- TF_LITE_ENSURE(context, NumDimensions(input0) < 4);
TF_LITE_ENSURE(context, NumDimensions(input0) >= data->axis);
// TODO(renjieliu): Support negative axis.
TF_LITE_ENSURE(context, data->axis >= 0);
diff --git a/tensorflow/lite/kernels/pack_test.cc b/tensorflow/lite/kernels/pack_test.cc
index 4f58deb..530cc2e 100644
--- a/tensorflow/lite/kernels/pack_test.cc
+++ b/tensorflow/lite/kernels/pack_test.cc
@@ -82,6 +82,19 @@
ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
}
+TEST(PackOpTest, FloatFiveDimensions) {
+ PackOpModel<float> model({TensorType_FLOAT32, {2, 2, 2, 2}}, 1, 2);
+ model.SetInput(0, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+ model.SetInput(
+ 1, {17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32});
+ model.Invoke();
+ EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 2, 2, 2, 2));
+ EXPECT_THAT(model.GetOutput(),
+ ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 17, 18, 19,
+ 20, 21, 22, 23, 24, 9, 10, 11, 12, 13, 14,
+ 15, 16, 25, 26, 27, 28, 29, 30, 31, 32}));
+}
+
// int32 tests.
TEST(PackOpTest, Int32ThreeInputs) {
PackOpModel<int32_t> model({TensorType_INT32, {2}}, 0, 3);
diff --git a/tensorflow/lite/kernels/register.cc b/tensorflow/lite/kernels/register.cc
index 3c60d28..c0e6f69 100644
--- a/tensorflow/lite/kernels/register.cc
+++ b/tensorflow/lite/kernels/register.cc
@@ -222,7 +222,9 @@
AddBuiltin(BuiltinOperator_LOG, Register_LOG());
AddBuiltin(BuiltinOperator_LOG_SOFTMAX, Register_LOG_SOFTMAX());
AddBuiltin(BuiltinOperator_CAST, Register_CAST());
- AddBuiltin(BuiltinOperator_DEQUANTIZE, Register_DEQUANTIZE());
+ AddBuiltin(BuiltinOperator_DEQUANTIZE, Register_DEQUANTIZE(),
+ /* min_version */ 1,
+ /* max_version */ 2);
AddBuiltin(BuiltinOperator_PRELU, Register_PRELU());
AddBuiltin(BuiltinOperator_MAXIMUM, Register_MAXIMUM());
AddBuiltin(BuiltinOperator_MINIMUM, Register_MINIMUM());
diff --git a/tensorflow/lite/kernels/resize_bilinear_test.cc b/tensorflow/lite/kernels/resize_bilinear_test.cc
index 530bb32..d3f4837 100644
--- a/tensorflow/lite/kernels/resize_bilinear_test.cc
+++ b/tensorflow/lite/kernels/resize_bilinear_test.cc
@@ -26,8 +26,8 @@
class ResizeBilinearOpModel : public SingleOpModel {
public:
- ResizeBilinearOpModel(const TensorData& input,
- std::initializer_list<int> size_data = {}) {
+ explicit ResizeBilinearOpModel(const TensorData& input,
+ std::initializer_list<int> size_data = {}) {
bool const_size = size_data.size() != 0;
input_ = AddInput(input);
if (const_size) {
diff --git a/tensorflow/lite/kernels/test_util.cc b/tensorflow/lite/kernels/test_util.cc
index 6b2a1f8..549ea78 100644
--- a/tensorflow/lite/kernels/test_util.cc
+++ b/tensorflow/lite/kernels/test_util.cc
@@ -129,14 +129,14 @@
interpreter_->SetAllowFp16PrecisionForFp32(allow_fp32_relax_to_fp16);
+ CHECK(interpreter_->AllocateTensors() == kTfLiteOk)
+ << "Cannot allocate tensors";
+ interpreter_->ResetVariableTensors();
+
// Modify delegate with function.
if (apply_delegate_fn_) {
apply_delegate_fn_(interpreter_.get());
}
-
- CHECK(interpreter_->AllocateTensors() == kTfLiteOk)
- << "Cannot allocate tensors";
- interpreter_->ResetVariableTensors();
}
void SingleOpModel::Invoke() { CHECK(interpreter_->Invoke() == kTfLiteOk); }
diff --git a/tensorflow/lite/kernels/test_util.h b/tensorflow/lite/kernels/test_util.h
index dadabb8..f5c67c3 100644
--- a/tensorflow/lite/kernels/test_util.h
+++ b/tensorflow/lite/kernels/test_util.h
@@ -307,10 +307,12 @@
if (is_quantized) {
if (t.min != 0 || t.max != 0) {
- // TODO(b/119422369): Handle signed int8 here.
if (t.type == TensorType_UINT8) {
std::tie(t.scale, t.zero_point) =
QuantizationParams<uint8_t>(t.min, t.max);
+ } else if (t.type == TensorType_INT8) {
+ std::tie(t.scale, t.zero_point) =
+ QuantizationParams<int8_t>(t.min, t.max);
} else if (t.type == TensorType_INT32) {
std::tie(t.scale, t.zero_point) =
QuantizationParams<int32_t>(t.min, t.max);
diff --git a/tensorflow/lite/lib_package/create_ios_frameworks.sh b/tensorflow/lite/lib_package/create_ios_frameworks.sh
index 7901655b..abf40e7 100755
--- a/tensorflow/lite/lib_package/create_ios_frameworks.sh
+++ b/tensorflow/lite/lib_package/create_ios_frameworks.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -x
+#!/bin/bash
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,10 +20,41 @@
echo "Starting"
TFLITE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.."
+usage() {
+ echo "Usage: $(basename "$0") [-a]"
+ echo "-g build with GPU delegate"
+ exit 1
+}
+
+USE_GPU_DELEGATE="false"
+FRAMEWORK_NAME="tensorflow_lite"
+while getopts "g" opt_name; do
+ case "$opt_name" in
+ g)
+ USE_GPU_DELEGATE="true"
+ FRAMEWORK_NAME="tensorflow_lite_gpu"
+ ;;
+ *) usage;;
+ esac
+done
+shift $((OPTIND - 1))
+readonly USE_GPU_DELEGATE
+readonly FRAMEWORK_NAME
+
+if [ $USE_GPU_DELEGATE == "true" ] ; then
+ for filename in metal_delegate.h libmetal_delegate.a ; do
+ if [[ ! -f "${TFLITE_DIR}/delegates/gpu/${filename}" ]] ; then
+ echo "File ${TFLITE_DIR}/delegates/gpu/${filename} doesn't exist."
+ echo "It's requried for building TFLite Framework with GPU. Aborting."
+ exit 1
+ fi
+ done
+fi
+
TMP_DIR=$(mktemp -d)
echo "Package dir: " $TMP_DIR
FW_DIR=$TMP_DIR/tensorflow_lite_ios_frameworks
-FW_DIR_TFLITE=$FW_DIR/tensorflow_lite.framework
+FW_DIR_TFLITE=$FW_DIR/$FRAMEWORK_NAME.framework
FW_DIR_TFLITE_HDRS=$FW_DIR_TFLITE/Headers
echo "Creating target Headers directories"
@@ -58,8 +89,14 @@
$FW_DIR_TFLITE
echo "Copying static libraries"
+# Note: There must be a static library with the same name
+# as the framework name.
cp $TFLITE_DIR/tools/make/gen/lib/libtensorflow-lite.a \
- $FW_DIR_TFLITE/tensorflow_lite
+ $FW_DIR_TFLITE/$FRAMEWORK_NAME
+if [ $USE_GPU_DELEGATE == "true" ] ; then
+ cp "${TFLITE_DIR}/delegates/gpu/libmetal_delegate.a" \
+ $FW_DIR_TFLITE/libmetal_delegate.a
+fi
# This is required, otherwise they interfere with the documentation of the
# pod at cocoapods.org.
@@ -71,10 +108,10 @@
TARGET_GEN_LOCATION="$TFLITE_DIR/gen/ios_frameworks"
echo "Moving results to target: " $TARGET_GEN_LOCATION
cd $FW_DIR
-zip -q -r tensorflow_lite.framework.zip tensorflow_lite.framework -x .DS_Store
+zip -q -r $FRAMEWORK_NAME.framework.zip $FRAMEWORK_NAME.framework -x .DS_Store
rm -rf $TARGET_GEN_LOCATION
mkdir -p $TARGET_GEN_LOCATION
-cp -r tensorflow_lite.framework.zip $TARGET_GEN_LOCATION
+cp -r $FRAMEWORK_NAME.framework.zip $TARGET_GEN_LOCATION
echo "Cleaning up"
rm -rf $TMP_DIR
diff --git a/tensorflow/lite/model.cc b/tensorflow/lite/model.cc
index 5ac0532..831c81a 100644
--- a/tensorflow/lite/model.cc
+++ b/tensorflow/lite/model.cc
@@ -121,6 +121,24 @@
return model;
}
+std::unique_ptr<FlatBufferModel> FlatBufferModel::VerifyAndBuildFromBuffer(
+ const char* buffer, size_t buffer_size, TfLiteVerifier* verifier,
+ ErrorReporter* error_reporter) {
+ error_reporter = ValidateErrorReporter(error_reporter);
+
+ flatbuffers::Verifier base_verifier(reinterpret_cast<const uint8_t*>(buffer),
+ buffer_size);
+ if (!VerifyModelBuffer(base_verifier)) {
+ return nullptr;
+ }
+
+ if (verifier && !verifier->Verify(buffer, buffer_size, error_reporter)) {
+ return nullptr;
+ }
+
+ return BuildFromBuffer(buffer, buffer_size, error_reporter);
+}
+
std::unique_ptr<FlatBufferModel> FlatBufferModel::BuildFromModel(
const tflite::Model* model_spec, ErrorReporter* error_reporter) {
error_reporter = ValidateErrorReporter(error_reporter);
diff --git a/tensorflow/lite/model.h b/tensorflow/lite/model.h
index 01e7c68..58c9767 100644
--- a/tensorflow/lite/model.h
+++ b/tensorflow/lite/model.h
@@ -80,10 +80,28 @@
// is destroyed. Caller retains ownership of `error_reporter` and must ensure
// its lifetime is longer than the FlatBufferModel instance.
// Returns a nullptr in case of failure.
+ // NOTE: this does NOT validate the buffer so it should NOT be called on
+ // invalid/untrusted input. Use VerifyAndBuildFromBuffer in that case
static std::unique_ptr<FlatBufferModel> BuildFromBuffer(
const char* buffer, size_t buffer_size,
ErrorReporter* error_reporter = DefaultErrorReporter());
+ // Verifies whether the content of the buffer is legit, then builds a model
+ // based on the pre-loaded flatbuffer.
+ // The verifier argument is an additional optional verifier for the buffer. By
+ // default, we always check with tflite::VerifyModelBuffer. If verifier is
+ // supplied, the buffer is checked against the verifier after the check
+ // against tflite::VerifyModelBuilder.
+ // The caller retains ownership of the buffer and should keep it alive until
+ // the returned object is destroyed. Caller retains ownership of
+ // `error_reporter` and must ensure its lifetime is longer than the
+ // FlatBufferModel instance.
+ // Returns a nullptr in case of failure.
+ static std::unique_ptr<FlatBufferModel> VerifyAndBuildFromBuffer(
+ const char* buffer, size_t buffer_size,
+ TfLiteVerifier* verifier = nullptr,
+ ErrorReporter* error_reporter = DefaultErrorReporter());
+
// Builds a model directly from a flatbuffer pointer. The caller retains
// ownership of the buffer and should keep it alive until the returned object
// is destroyed. Caller retains ownership of `error_reporter` and must ensure
diff --git a/tensorflow/lite/models/smartreply/demo/app/src/main/BUILD b/tensorflow/lite/models/smartreply/demo/app/src/main/BUILD
index b14af4c..73326e9 100644
--- a/tensorflow/lite/models/smartreply/demo/app/src/main/BUILD
+++ b/tensorflow/lite/models/smartreply/demo/app/src/main/BUILD
@@ -62,6 +62,7 @@
],
deps = [
"//tensorflow/lite:framework",
+ "//tensorflow/lite/java/jni",
"//tensorflow/lite/models/smartreply:predictor_lib",
],
alwayslink = 1,
diff --git a/tensorflow/lite/profiling/BUILD b/tensorflow/lite/profiling/BUILD
index c7a8e4f..52ea6fe 100644
--- a/tensorflow/lite/profiling/BUILD
+++ b/tensorflow/lite/profiling/BUILD
@@ -58,7 +58,6 @@
name = "profile_summarizer_test",
srcs = ["profile_summarizer_test.cc"],
copts = common_copts,
- tags = ["no_oss"],
deps = [
":profile_summarizer",
"//tensorflow/lite:framework",
diff --git a/tensorflow/lite/profiling/profiler_test.cc b/tensorflow/lite/profiling/profiler_test.cc
index 82d0537..addebab 100644
--- a/tensorflow/lite/profiling/profiler_test.cc
+++ b/tensorflow/lite/profiling/profiler_test.cc
@@ -27,11 +27,8 @@
namespace profiling {
namespace {
-void AssertDurationOfEventAroundMs(const ProfileEvent* event,
- double expected_ms, double eps_ms) {
- double duration_ms =
- (event->end_timestamp_us - event->begin_timestamp_us) / 1e3;
- EXPECT_NEAR(expected_ms, duration_ms, eps_ms);
+double GetDurationOfEventMs(const ProfileEvent* event) {
+ return (event->end_timestamp_us - event->begin_timestamp_us) / 1e3;
}
void SleepForQuarterSecond(Profiler* profiler) {
@@ -84,12 +81,17 @@
#ifndef ADDRESS_SANITIZER
// ASAN build is sometimes very slow. Set a large epsilon to avoid flakiness.
+ // Due to flakiness, just verify relative values match.
const int eps_ms = 50;
- AssertDurationOfEventAroundMs(profile_events[0], /*expected_ms*/ 500, eps_ms);
- AssertDurationOfEventAroundMs(profile_events[1], /*expected_ms*/ 250, eps_ms);
- AssertDurationOfEventAroundMs(profile_events[2], /*expected_ms*/ 250, eps_ms);
- AssertDurationOfEventAroundMs(profile_events[3], /*expected_ms*/ 250, eps_ms);
- AssertDurationOfEventAroundMs(profile_events[4], /*expected_ms*/ 250, eps_ms);
+ auto parent_ms = GetDurationOfEventMs(profile_events[0]);
+ double child_ms[2], sleep_for_quarter_ms[2];
+ child_ms[0] = GetDurationOfEventMs(profile_events[1]);
+ child_ms[1] = GetDurationOfEventMs(profile_events[3]);
+ sleep_for_quarter_ms[0] = GetDurationOfEventMs(profile_events[2]);
+ sleep_for_quarter_ms[1] = GetDurationOfEventMs(profile_events[4]);
+ EXPECT_NEAR(parent_ms, child_ms[0] + child_ms[1], eps_ms);
+ EXPECT_NEAR(child_ms[0], sleep_for_quarter_ms[0], eps_ms);
+ EXPECT_NEAR(child_ms[1], sleep_for_quarter_ms[1], eps_ms);
#endif
}
diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py
index 563312e..9c60399 100644
--- a/tensorflow/lite/python/convert.py
+++ b/tensorflow/lite/python/convert.py
@@ -97,6 +97,7 @@
return result
+@_tf_export("lite.OpsSet")
class OpsSet(enum.Enum):
"""Enum class defining the sets of ops available to generate TFLite models.
diff --git a/tensorflow/lite/python/convert_saved_model_test.py b/tensorflow/lite/python/convert_saved_model_test.py
index 7611385..fdcbc79 100644
--- a/tensorflow/lite/python/convert_saved_model_test.py
+++ b/tensorflow/lite/python/convert_saved_model_test.py
@@ -39,6 +39,7 @@
class TensorFunctionsTest(test_util.TensorFlowTestCase):
+ @test_util.run_v1_only("b/120545219")
def testGetTensorsValid(self):
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32)
@@ -49,6 +50,7 @@
sess.graph, ["Placeholder"])
self.assertEqual("Placeholder:0", tensors[0].name)
+ @test_util.run_v1_only("b/120545219")
def testGetTensorsInvalid(self):
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32)
@@ -61,6 +63,7 @@
self.assertEqual("Invalid tensors 'invalid-input' were found.",
str(error.exception))
+ @test_util.run_v1_only("b/120545219")
def testSetTensorShapeValid(self):
tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32)
self.assertEqual([None, 3, 5], tensor.shape.as_list())
@@ -68,6 +71,7 @@
convert_saved_model.set_tensor_shapes([tensor], {"Placeholder": [5, 3, 5]})
self.assertEqual([5, 3, 5], tensor.shape.as_list())
+ @test_util.run_v1_only("b/120545219")
def testSetTensorShapeNoneValid(self):
tensor = array_ops.placeholder(dtype=dtypes.float32)
self.assertEqual(None, tensor.shape)
@@ -75,6 +79,7 @@
convert_saved_model.set_tensor_shapes([tensor], {"Placeholder": [1, 3, 5]})
self.assertEqual([1, 3, 5], tensor.shape.as_list())
+ @test_util.run_v1_only("b/120545219")
def testSetTensorShapeArrayInvalid(self):
# Tests set_tensor_shape where the tensor name passed in doesn't exist.
tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32)
@@ -88,6 +93,7 @@
str(error.exception))
self.assertEqual([None, 3, 5], tensor.shape.as_list())
+ @test_util.run_deprecated_v1
def testSetTensorShapeDimensionInvalid(self):
# Tests set_tensor_shape where the shape passed in is incompatiable.
tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32)
@@ -96,11 +102,11 @@
with self.assertRaises(ValueError) as error:
convert_saved_model.set_tensor_shapes([tensor],
{"Placeholder": [1, 5, 5]})
- self.assertIn(
- "The shape of tensor 'Placeholder' cannot be changed from "
- "(?, 3, 5) to [1, 5, 5].", str(error.exception))
+ self.assertIn("The shape of tensor 'Placeholder' cannot be changed",
+ str(error.exception))
self.assertEqual([None, 3, 5], tensor.shape.as_list())
+ @test_util.run_v1_only("b/120545219")
def testSetTensorShapeEmpty(self):
tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32)
self.assertEqual([None, 3, 5], tensor.shape.as_list())
diff --git a/tensorflow/lite/python/convert_test.py b/tensorflow/lite/python/convert_test.py
index 2a6f1f6..cf49ee2 100644
--- a/tensorflow/lite/python/convert_test.py
+++ b/tensorflow/lite/python/convert_test.py
@@ -34,6 +34,7 @@
from tensorflow.python.platform import test
+@test_util.run_v1_only("b/120545219")
class ConvertTest(test_util.TensorFlowTestCase):
def testBasic(self):
@@ -176,6 +177,7 @@
"QUANTIZED_UINT8.", str(error.exception))
+@test_util.run_v1_only("b/120545219")
class ConvertTestOpHint(test_util.TensorFlowTestCase):
"""Test the hint to stub functionality."""
diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py
index 1ae0d3c..1f9c768 100644
--- a/tensorflow/lite/python/lite_test.py
+++ b/tensorflow/lite/python/lite_test.py
@@ -80,6 +80,7 @@
self.assertTrue(converter._has_valid_tensors())
+@test_util.run_v1_only('b/120545219')
class FromSessionTest(test_util.TensorFlowTestCase):
def testFloat(self):
@@ -177,6 +178,38 @@
'Quantization input stats are not available for input tensors '
'\'inputB\'.', str(error.exception))
+ def testIntermediateInputArray(self):
+ """Convert a model from an intermediate input array."""
+ in_tensor_init = array_ops.placeholder(
+ shape=[1, 16, 16, 3], dtype=dtypes.float32)
+ in_tensor_final = in_tensor_init + in_tensor_init
+ out_tensor = in_tensor_final + in_tensor_final
+ sess = session.Session()
+
+ # Convert model and ensure model is not None.
+ converter = lite.TFLiteConverter.from_session(sess, [in_tensor_final],
+ [out_tensor])
+ tflite_model = converter.convert()
+ self.assertTrue(tflite_model)
+
+ # Check values from converted model.
+ interpreter = Interpreter(model_content=tflite_model)
+ interpreter.allocate_tensors()
+
+ input_details = interpreter.get_input_details()
+ self.assertEqual(1, len(input_details))
+ self.assertEqual('add', input_details[0]['name'])
+ self.assertEqual(np.float32, input_details[0]['dtype'])
+ self.assertTrue(([1, 16, 16, 3] == input_details[0]['shape']).all())
+ self.assertEqual((0., 0.), input_details[0]['quantization'])
+
+ output_details = interpreter.get_output_details()
+ self.assertEqual(1, len(output_details))
+ self.assertEqual('add_1', output_details[0]['name'])
+ self.assertEqual(np.float32, output_details[0]['dtype'])
+ self.assertTrue(([1, 16, 16, 3] == output_details[0]['shape']).all())
+ self.assertEqual((0., 0.), output_details[0]['quantization'])
+
def testSizeNoneInvalid(self):
in_tensor = array_ops.placeholder(dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
@@ -465,6 +498,7 @@
interpreter.allocate_tensors()
+@test_util.run_v1_only('b/120545219')
class FromFrozenGraphFile(test_util.TensorFlowTestCase):
def testFloat(self):
@@ -712,6 +746,7 @@
interpreter.allocate_tensors()
+@test_util.run_v1_only('b/120545219')
class FromSavedModelTest(test_util.TensorFlowTestCase):
def _createSavedModel(self, shape):
@@ -856,6 +891,7 @@
interpreter.allocate_tensors()
+@test_util.run_v1_only('b/120545219')
class FromKerasFile(test_util.TensorFlowTestCase):
def setUp(self):
diff --git a/tensorflow/lite/schema/builtin_ops_header/BUILD b/tensorflow/lite/schema/builtin_ops_header/BUILD
index 8a01541..52cbd05 100644
--- a/tensorflow/lite/schema/builtin_ops_header/BUILD
+++ b/tensorflow/lite/schema/builtin_ops_header/BUILD
@@ -24,7 +24,6 @@
cc_test(
name = "generator_test",
srcs = ["generator_test.cc"],
- tags = ["no_oss"],
deps = [
":generator",
"@com_google_googletest//:gtest",
@@ -37,7 +36,6 @@
data = [
"//tensorflow/lite:builtin_ops.h",
],
- tags = ["no_oss"],
deps = [
":generator",
"@com_google_googletest//:gtest",
diff --git a/tensorflow/lite/schema/schema.fbs b/tensorflow/lite/schema/schema.fbs
index 6436167..980f13b 100644
--- a/tensorflow/lite/schema/schema.fbs
+++ b/tensorflow/lite/schema/schema.fbs
@@ -45,7 +45,7 @@
// Custom quantization parameters for experimenting with new quantization
// techniques.
table CustomQuantization {
- custom:[byte];
+ custom:[ubyte] (force_align: 16);
}
// Represents a specific quantization technique's parameters.
diff --git a/tensorflow/lite/schema/schema_generated.h b/tensorflow/lite/schema/schema_generated.h
index af8b143..637cbaf 100755
--- a/tensorflow/lite/schema/schema_generated.h
+++ b/tensorflow/lite/schema/schema_generated.h
@@ -2247,7 +2247,7 @@
struct CustomQuantizationT : public flatbuffers::NativeTable {
typedef CustomQuantization TableType;
- std::vector<int8_t> custom;
+ std::vector<uint8_t> custom;
CustomQuantizationT() {
}
};
@@ -2257,8 +2257,8 @@
enum {
VT_CUSTOM = 4
};
- const flatbuffers::Vector<int8_t> *custom() const {
- return GetPointer<const flatbuffers::Vector<int8_t> *>(VT_CUSTOM);
+ const flatbuffers::Vector<uint8_t> *custom() const {
+ return GetPointer<const flatbuffers::Vector<uint8_t> *>(VT_CUSTOM);
}
bool Verify(flatbuffers::Verifier &verifier) const {
return VerifyTableStart(verifier) &&
@@ -2274,7 +2274,7 @@
struct CustomQuantizationBuilder {
flatbuffers::FlatBufferBuilder &fbb_;
flatbuffers::uoffset_t start_;
- void add_custom(flatbuffers::Offset<flatbuffers::Vector<int8_t>> custom) {
+ void add_custom(flatbuffers::Offset<flatbuffers::Vector<uint8_t>> custom) {
fbb_.AddOffset(CustomQuantization::VT_CUSTOM, custom);
}
explicit CustomQuantizationBuilder(flatbuffers::FlatBufferBuilder &_fbb)
@@ -2291,7 +2291,7 @@
inline flatbuffers::Offset<CustomQuantization> CreateCustomQuantization(
flatbuffers::FlatBufferBuilder &_fbb,
- flatbuffers::Offset<flatbuffers::Vector<int8_t>> custom = 0) {
+ flatbuffers::Offset<flatbuffers::Vector<uint8_t>> custom = 0) {
CustomQuantizationBuilder builder_(_fbb);
builder_.add_custom(custom);
return builder_.Finish();
@@ -2299,10 +2299,10 @@
inline flatbuffers::Offset<CustomQuantization> CreateCustomQuantizationDirect(
flatbuffers::FlatBufferBuilder &_fbb,
- const std::vector<int8_t> *custom = nullptr) {
+ const std::vector<uint8_t> *custom = nullptr) {
return tflite::CreateCustomQuantization(
_fbb,
- custom ? _fbb.CreateVector<int8_t>(*custom) : 0);
+ custom ? _fbb.CreateVector<uint8_t>(*custom) : 0);
}
flatbuffers::Offset<CustomQuantization> CreateCustomQuantization(flatbuffers::FlatBufferBuilder &_fbb, const CustomQuantizationT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
diff --git a/tensorflow/lite/testing/generated_examples_zip_test.cc b/tensorflow/lite/testing/generated_examples_zip_test.cc
index 91a4851..a9a31ad 100644
--- a/tensorflow/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/lite/testing/generated_examples_zip_test.cc
@@ -205,7 +205,7 @@
}
if (!added) {
string message = "Test had no examples: " + original_file;
- return tensorflow::Status(tensorflow::error::UNKNOWN, message.c_str());
+ return tensorflow::Status(tensorflow::error::UNKNOWN, message);
}
return tensorflow::Status::OK();
}
diff --git a/tensorflow/lite/testing/join.h b/tensorflow/lite/testing/join.h
index 7d0040c..d1c3146 100644
--- a/tensorflow/lite/testing/join.h
+++ b/tensorflow/lite/testing/join.h
@@ -16,6 +16,7 @@
#define TENSORFLOW_LITE_TESTING_JOIN_H_
#include <cstdlib>
+#include <iomanip>
#include <sstream>
#include "tensorflow/lite/string.h"
@@ -30,9 +31,9 @@
return "";
}
std::stringstream result;
- result << data[0];
+ result << std::setprecision(9) << data[0];
for (int i = 1; i < len; i++) {
- result << delimiter << data[i];
+ result << std::setprecision(9) << delimiter << data[i];
}
return result.str();
}
diff --git a/tensorflow/lite/testing/join_test.cc b/tensorflow/lite/testing/join_test.cc
index a8d036c..0b3c07f 100644
--- a/tensorflow/lite/testing/join_test.cc
+++ b/tensorflow/lite/testing/join_test.cc
@@ -28,7 +28,7 @@
TEST(JoinTest, JoinFloat) {
float data[] = {1.0, -3, 2.3, 1e-5};
- EXPECT_EQ(Join(data, 4, " "), "1 -3 2.3 1e-05");
+ EXPECT_EQ(Join(data, 4, " "), "1 -3 2.29999995 9.99999975e-06");
}
TEST(JoinTest, JoinNullData) { EXPECT_THAT(Join<int>(nullptr, 3, ","), ""); }
diff --git a/tensorflow/lite/testing/model_coverage/model_coverage_lib_test.py b/tensorflow/lite/testing/model_coverage/model_coverage_lib_test.py
index 6b4e742..4e329ac 100644
--- a/tensorflow/lite/testing/model_coverage/model_coverage_lib_test.py
+++ b/tensorflow/lite/testing/model_coverage/model_coverage_lib_test.py
@@ -28,6 +28,7 @@
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
@@ -42,6 +43,7 @@
write_graph(sess.graph_def, '', graph_def_file, False)
return graph_def_file
+ @test_util.run_v1_only('b/120545219')
def testFloat(self):
with session.Session().as_default() as sess:
in_tensor = array_ops.placeholder(
@@ -51,6 +53,7 @@
model_coverage.test_frozen_graph(filename, ['Placeholder'], ['add'])
+ @test_util.run_v1_only('b/120545219')
def testMultipleOutputs(self):
with session.Session().as_default() as sess:
in_tensor_1 = array_ops.placeholder(
@@ -84,15 +87,18 @@
filename = self._saveFrozenGraph(sess)
return filename
+ @test_util.run_v1_only('b/120545219')
def testQuantized(self):
filename = self._getQuantizedModel()
model_coverage.test_frozen_graph_quant(filename, ['inputA'], ['output'])
+ @test_util.run_v1_only('b/120545219')
def testQuantizedInputShapes(self):
filename = self._getQuantizedModel()
model_coverage.test_frozen_graph_quant(
filename, ['inputA'], ['output'], input_shapes={'inputA': [33, 33]})
+ @test_util.run_v1_only('b/120545219')
def testQuantizedFlexAll(self):
filename = self._getQuantizedModel()
model_coverage.test_frozen_graph_quant(
@@ -102,6 +108,7 @@
class EvaluateSavedModel(test.TestCase):
+ @test_util.run_v1_only('b/120545219')
def testFloat(self):
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
with session.Session().as_default() as sess:
@@ -139,18 +146,21 @@
os.close(fd)
return keras_file
+ @test_util.run_v1_only('b/120545219')
def testFloat(self):
model = self._getSingleInputKerasModel()
keras_file = self._saveKerasModel(model)
model_coverage.test_keras_model(keras_file)
+ @test_util.run_v1_only('b/120545219')
def testPostTrainingQuantize(self):
model = self._getSingleInputKerasModel()
keras_file = self._saveKerasModel(model)
model_coverage.test_keras_model(keras_file, post_training_quantize=True)
+ @test_util.run_v1_only('b/120545219')
def testTargetOps(self):
model = self._getSingleInputKerasModel()
keras_file = self._saveKerasModel(model)
diff --git a/tensorflow/lite/testing/tf_driver_test.cc b/tensorflow/lite/testing/tf_driver_test.cc
index 4381fe4..363d162 100644
--- a/tensorflow/lite/testing/tf_driver_test.cc
+++ b/tensorflow/lite/testing/tf_driver_test.cc
@@ -46,7 +46,7 @@
TestDriver driver;
ASSERT_EQ(driver.WriteAndReadBack(tensorflow::DT_FLOAT, {1, 2, 2},
"0.10,0.20,0.30,0.40"),
- "0.1,0.2,0.3,0.4");
+ "0.100000001,0.200000003,0.300000012,0.400000006");
ASSERT_EQ(driver.WriteAndReadBack(tensorflow::DT_INT32, {1, 2, 2},
"10,40,100,-100"),
"10,40,100,-100");
@@ -111,8 +111,10 @@
runner->ResetTensor(2);
runner->Invoke();
- ASSERT_EQ(runner->ReadOutput(0), "0.101,0.202,0.303,0.404");
- ASSERT_EQ(runner->ReadOutput(1), "0.011,0.022,0.033,0.044");
+ ASSERT_EQ(runner->ReadOutput(0),
+ "0.101000004,0.202000007,0.303000003,0.404000014");
+ ASSERT_EQ(runner->ReadOutput(1),
+ "0.0109999999,0.0219999999,0.0329999998,0.0439999998");
}
} // namespace
diff --git a/tensorflow/lite/testing/tflite_driver.cc b/tensorflow/lite/testing/tflite_driver.cc
index 27e3a37..4e11d49 100644
--- a/tensorflow/lite/testing/tflite_driver.cc
+++ b/tensorflow/lite/testing/tflite_driver.cc
@@ -77,6 +77,13 @@
SetTensorData(values, &data_);
}
+ template <>
+ void SetData<string>(const string& csv_values) {
+ string s = absl::HexStringToBytes(csv_values);
+ data_.raw = new char[s.size()];
+ memcpy(data_.raw, s.data(), s.size());
+ }
+
bool Check(bool verbose, const TfLiteTensor& tensor) {
switch (tensor.type) {
case kTfLiteFloat32:
@@ -89,6 +96,8 @@
return TypedCheck<uint8_t>(verbose, tensor);
case kTfLiteBool:
return TypedCheck<bool>(verbose, tensor);
+ case kTfLiteString:
+ return TypedCheck<string>(verbose, tensor);
default:
fprintf(stderr, "Unsupported type %d in Check\n", tensor.type);
return false;
@@ -135,6 +144,46 @@
return good_output;
}
+ template <>
+ bool TypedCheck<string>(bool verbose, const TfLiteTensor& tensor) {
+ if (tensor.data.raw == nullptr) {
+ if (verbose) {
+ std::cerr << " got empty string" << std::endl;
+ }
+ return false;
+ }
+ int expected_num_strings = GetStringCount(data_.raw);
+ int returned_num_strings = GetStringCount(tensor.data.raw);
+ if (expected_num_strings != returned_num_strings) {
+ if (verbose) {
+ std::cerr << " string count differ: got " << returned_num_strings
+ << ", but expected " << expected_num_strings << std::endl;
+ }
+ return false;
+ }
+ for (int i = 0; i < returned_num_strings; ++i) {
+ auto expected_ref = GetString(data_.raw, i);
+ auto returned_ref = GetString(tensor.data.raw, i);
+ if (expected_ref.len != returned_ref.len) {
+ if (verbose) {
+ std::cerr << " index " << i << ": got string of size "
+ << returned_ref.len << ", but expected size "
+ << expected_ref.len << std::endl;
+ }
+ return false;
+ }
+ if (strncmp(expected_ref.str, returned_ref.str, returned_ref.len) != 0) {
+ if (verbose) {
+ std::cerr << " index " << i << ": strings are different"
+ << std::endl;
+ }
+ return false;
+ }
+ }
+
+ return true;
+ }
+
TfLitePtrUnion data_;
size_t num_elements_;
};
@@ -250,8 +299,9 @@
break;
}
default:
- fprintf(stderr, "Unsupported type %d in SetInput\n", tensor->type);
- Invalidate("Unsupported tensor data type");
+ Invalidate(absl::StrCat("Unsupported tensor type ",
+ TfLiteTypeGetName(tensor->type),
+ " in TfLiteDriver::SetInput"));
return;
}
}
@@ -260,8 +310,7 @@
if (!IsValid()) return;
auto* tensor = interpreter_->tensor(id);
if (expected_output_.count(id) != 0) {
- fprintf(stderr, "Overridden expectation for tensor %d\n", id);
- Invalidate("Overridden expectation");
+ Invalidate(absl::StrCat("Overridden expectation for tensor '", id, "'"));
}
expected_output_[id].reset(new Expectation);
switch (tensor->type) {
@@ -280,9 +329,13 @@
case kTfLiteBool:
expected_output_[id]->SetData<bool>(csv_values);
break;
+ case kTfLiteString:
+ expected_output_[id]->SetData<string>(csv_values);
+ break;
default:
- fprintf(stderr, "Unsupported type %d in SetExpectation\n", tensor->type);
- Invalidate("Unsupported tensor data type");
+ Invalidate(absl::StrCat("Unsupported tensor type ",
+ TfLiteTypeGetName(tensor->type),
+ " in TfLiteDriver::SetExpectation"));
return;
}
}
diff --git a/tensorflow/lite/toco/BUILD b/tensorflow/lite/toco/BUILD
index 82aa1f5..93d41fc 100644
--- a/tensorflow/lite/toco/BUILD
+++ b/tensorflow/lite/toco/BUILD
@@ -341,7 +341,6 @@
tf_cc_test(
name = "import_tensorflow_test",
srcs = ["import_tensorflow_test.cc"],
- tags = ["no_oss"],
deps = [
":toco_tooling",
"//tensorflow/core:framework",
@@ -384,7 +383,6 @@
tf_cc_test(
name = "tooling_util_test",
srcs = ["tooling_util_test.cc"],
- tags = ["no_oss"],
deps = [
":model",
":tooling_util",
@@ -468,7 +466,6 @@
data = [
"toco_port_test.cc",
],
- tags = ["no_oss"],
deps = [
":toco_port",
"@com_google_googletest//:gtest_main",
diff --git a/tensorflow/lite/toco/graph_transformations/graph_transformations.h b/tensorflow/lite/toco/graph_transformations/graph_transformations.h
index 73a90c8..187b584 100644
--- a/tensorflow/lite/toco/graph_transformations/graph_transformations.h
+++ b/tensorflow/lite/toco/graph_transformations/graph_transformations.h
@@ -139,7 +139,7 @@
DECLARE_GRAPH_TRANSFORMATION(MoveBinaryOperatorBeforeReshape)
DECLARE_GRAPH_TRANSFORMATION(PropagateActivationFunctionIntoConstants)
DECLARE_GRAPH_TRANSFORMATION(PropagateArrayDataTypes)
-DECLARE_GRAPH_TRANSFORMATION(PropagateFakeQuantNumBits);
+DECLARE_GRAPH_TRANSFORMATION(PropagateFakeQuantNumBits)
DECLARE_GRAPH_TRANSFORMATION(PropagateFixedSizes)
DECLARE_GRAPH_TRANSFORMATION(HardcodeMinMax)
DECLARE_GRAPH_TRANSFORMATION(Quantize)
diff --git a/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc b/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc
index df50f31..2e41767 100644
--- a/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc
+++ b/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc
@@ -208,12 +208,32 @@
if (output_array.minmax) {
return false;
}
- const auto& input_array_1 = model->GetArray(op->inputs[1]);
- if (!input_array_1.minmax) {
+
+ auto& input_array_1 = model->GetArray(op->inputs[1]);
+ auto& input_array_2 = model->GetArray(op->inputs[2]);
+
+ if (!input_array_1.minmax && !input_array_2.minmax) {
return false;
}
- const auto& input_array_2 = model->GetArray(op->inputs[2]);
- if (!input_array_2.minmax) {
+
+ // Propagate up if one input is quantized and the other is constant.
+ if (!input_array_1.minmax &&
+ IsConstantParameterArray(*model, op->inputs[1])) {
+ auto& minmax_1 = input_array_1.GetOrCreateMinMax();
+ const auto& minmax_2 = input_array_2.GetMinMax();
+ minmax_1.min = minmax_2.min;
+ minmax_1.max = minmax_2.max;
+ }
+
+ if (!input_array_2.minmax &&
+ IsConstantParameterArray(*model, op->inputs[2])) {
+ auto& minmax_2 = input_array_2.GetOrCreateMinMax();
+ const auto& minmax_1 = input_array_1.GetMinMax();
+ minmax_2.min = minmax_1.min;
+ minmax_2.max = minmax_1.max;
+ }
+
+ if (!input_array_1.minmax || !input_array_2.minmax) {
return false;
}
diff --git a/tensorflow/lite/toco/graph_transformations/quantize.cc b/tensorflow/lite/toco/graph_transformations/quantize.cc
index 1146078..2fa80f2 100644
--- a/tensorflow/lite/toco/graph_transformations/quantize.cc
+++ b/tensorflow/lite/toco/graph_transformations/quantize.cc
@@ -64,6 +64,7 @@
type == OperatorType::kRelu1 || type == OperatorType::kRelu6 ||
type == OperatorType::kShape || type == OperatorType::kExpandDims ||
type == OperatorType::kPack || type == OperatorType::kTopK_V2 ||
+ type == OperatorType::kRandomUniform ||
type == OperatorType::kResizeNearestNeighbor ||
type == OperatorType::kPRelu;
}
diff --git a/tensorflow/lite/toco/graph_transformations/tests/BUILD b/tensorflow/lite/toco/graph_transformations/tests/BUILD
index 2e9b213..bbbedbe 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/BUILD
+++ b/tensorflow/lite/toco/graph_transformations/tests/BUILD
@@ -10,7 +10,6 @@
tf_cc_test(
name = "lstm_utils_test",
srcs = ["lstm_utils_test.cc"],
- tags = ["no_oss"],
deps = [
"//tensorflow/lite/toco:graph_transformations",
"//tensorflow/lite/toco:model",
@@ -22,7 +21,6 @@
tf_cc_test(
name = "resolve_constant_concatenation_test",
srcs = ["resolve_constant_concatenation_test.cc"],
- tags = ["no_oss"],
deps = [
"//tensorflow/lite/toco:graph_transformations",
"//tensorflow/lite/toco:model",
@@ -34,7 +32,6 @@
tf_cc_test(
name = "resolve_constant_unary_test",
srcs = ["resolve_constant_unary_test.cc"],
- tags = ["no_oss"],
deps = [
"//tensorflow/lite/toco:graph_transformations",
"//tensorflow/lite/toco:model",
diff --git a/tensorflow/lite/toco/python/BUILD b/tensorflow/lite/toco/python/BUILD
index 07056f6..8a6e82e 100644
--- a/tensorflow/lite/toco/python/BUILD
+++ b/tensorflow/lite/toco/python/BUILD
@@ -1,4 +1,8 @@
-package(default_visibility = ["//visibility:public"])
+package(default_visibility = [
+ "//tensorflow/contrib/lite:__subpackages__",
+ "//tensorflow/lite:__subpackages__",
+ "//tensorflow/tools/pip_package:__subpackages__",
+])
licenses(["notice"]) # Apache 2.0
@@ -9,7 +13,10 @@
config_setting(
name = "tflite_convert_with_select_tf_ops",
define_values = {"tflite_convert_with_select_tf_ops": "true"},
- visibility = ["//visibility:public"],
+ visibility = [
+ "//tensorflow/contrib/lite:__subpackages__",
+ "//tensorflow/lite:__subpackages__",
+ ],
)
cc_library(
@@ -37,6 +44,12 @@
tf_py_wrap_cc(
name = "tensorflow_wrap_toco",
srcs = ["toco.i"],
+ visibility = [
+ "//learning/expander/pod/deep_pod/utils:__subpackages__",
+ "//research/handwriting/converters/tflite:__subpackages__",
+ "//tensorflow/contrib/lite:__subpackages__",
+ "//tensorflow/lite:__subpackages__",
+ ],
deps = [
":toco_python_api",
"//tensorflow/lite/toco:model_flags_proto_cc",
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/BUILD b/tensorflow/lite/toco/tensorflow_graph_matching/BUILD
index 56acc28..ae361bf 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/BUILD
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/BUILD
@@ -60,7 +60,6 @@
tf_cc_test(
name = "resolve_svdf_test",
srcs = ["resolve_svdf_test.cc"],
- tags = ["no_oss"],
deps = [
":cluster",
":cluster_utils",
diff --git a/tensorflow/lite/toco/tflite/BUILD b/tensorflow/lite/toco/tflite/BUILD
index 99c4f8e..36ca638 100644
--- a/tensorflow/lite/toco/tflite/BUILD
+++ b/tensorflow/lite/toco/tflite/BUILD
@@ -41,7 +41,6 @@
srcs = [
"operator_test.cc",
],
- tags = ["no_oss"],
deps = [
":operator",
"//tensorflow/core:ops",
@@ -72,7 +71,6 @@
srcs = [
"types_test.cc",
],
- tags = ["no_oss"],
deps = [
":types",
"//tensorflow/core:ops",
@@ -107,7 +105,6 @@
srcs = [
"export_test.cc",
],
- tags = ["no_oss"],
deps = [
":export",
"//tensorflow/core:ops",
@@ -142,7 +139,6 @@
srcs = [
"import_test.cc",
],
- tags = ["no_oss"],
deps = [
":import",
"//tensorflow/core:ops",
diff --git a/tensorflow/lite/toco/tflite/import.cc b/tensorflow/lite/toco/tflite/import.cc
index 88028aa..1692f72 100644
--- a/tensorflow/lite/toco/tflite/import.cc
+++ b/tensorflow/lite/toco/tflite/import.cc
@@ -165,21 +165,28 @@
}
}
-void ImportIOTensors(const ::tflite::Model& input_model,
+void ImportIOTensors(const ModelFlags& model_flags,
+ const ::tflite::Model& input_model,
const details::TensorsTable& tensors_table, Model* model) {
- auto inputs = (*input_model.subgraphs())[0]->inputs();
- if (inputs) {
- for (int input : *inputs) {
- const string& input_name = tensors_table.at(input);
- model->flags.add_input_arrays()->set_name(input_name);
+ // Import from the first subgraph if input arrays have not been specified.
+ if (model_flags.input_arrays().empty()) {
+ auto inputs = (*input_model.subgraphs())[0]->inputs();
+ if (inputs) {
+ for (int input : *inputs) {
+ const string& input_name = tensors_table.at(input);
+ model->flags.add_input_arrays()->set_name(input_name);
+ }
}
}
- auto outputs = (*input_model.subgraphs())[0]->outputs();
- if (outputs) {
- for (int output : *outputs) {
- const string& output_name = tensors_table.at(output);
- model->flags.add_output_arrays(output_name);
+ // Import from the first subgraph if output arrays have not been specified.
+ if (model_flags.output_arrays().empty()) {
+ auto outputs = (*input_model.subgraphs())[0]->outputs();
+ if (outputs) {
+ for (int output : *outputs) {
+ const string& output_name = tensors_table.at(output);
+ model->flags.add_output_arrays(output_name);
+ }
}
}
}
@@ -219,7 +226,8 @@
ImportTensors(*input_model, model.get());
ImportOperators(*input_model, ops_by_name, tensors_table, operators_table,
model.get());
- ImportIOTensors(*input_model, tensors_table, model.get());
+
+ ImportIOTensors(model_flags, *input_model, tensors_table, model.get());
UndoWeightsShuffling(model.get());
diff --git a/tensorflow/lite/toco/tflite/operator.cc b/tensorflow/lite/toco/tflite/operator.cc
index e0faed4..205af23 100644
--- a/tensorflow/lite/toco/tflite/operator.cc
+++ b/tensorflow/lite/toco/tflite/operator.cc
@@ -1482,6 +1482,30 @@
const bool enable_select_tf_ops_;
};
+class Dequantize
+ : public BuiltinOperator<DequantizeOperator, ::tflite::DequantizeOptions,
+ ::tflite::BuiltinOptions_DequantizeOptions> {
+ public:
+ using BuiltinOperator::BuiltinOperator;
+
+ flatbuffers::Offset<TfLiteOptions> WriteOptions(
+ const TocoOperator& op,
+ flatbuffers::FlatBufferBuilder* builder) const override {
+ return ::tflite::CreateDequantizeOptions(*builder);
+ }
+
+ void ReadOptions(const TfLiteOptions& options,
+ TocoOperator* op) const override {}
+
+ int GetVersion(const Operator& op) const override {
+ // TODO(suharshs): Dequantize now supports INT8 in addition to
+ // QUANTIZED_UINT8. When TOCO can create models with INT8, we need
+ // to find a way to see the type here and return version 2. Right now
+ // version 2 will only be added by post training quantization tools.
+ return 1;
+ }
+};
+
namespace {
// Build a vector containing all the known operators.
std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList(
diff --git a/tensorflow/lite/tools/BUILD b/tensorflow/lite/tools/BUILD
index 93725b5..1d141b5 100644
--- a/tensorflow/lite/tools/BUILD
+++ b/tensorflow/lite/tools/BUILD
@@ -53,7 +53,6 @@
"//tensorflow/lite:testdata/test_model_broken.bin",
],
tags = [
- "no_oss",
"tflite_not_portable_android",
"tflite_not_portable_ios",
],
@@ -80,7 +79,6 @@
size = "small",
srcs = ["verifier_test.cc"],
tags = [
- "no_oss",
"tflite_not_portable",
],
deps = [
diff --git a/tensorflow/lite/tools/benchmark/BUILD b/tensorflow/lite/tools/benchmark/BUILD
index 583046a..bc47406 100644
--- a/tensorflow/lite/tools/benchmark/BUILD
+++ b/tensorflow/lite/tools/benchmark/BUILD
@@ -70,6 +70,7 @@
deps = [
":benchmark_tflite_model_lib",
":command_line_flags",
+ "//tensorflow/lite:framework",
"//tensorflow/lite/testing:util",
"@com_google_googletest//:gtest",
],
diff --git a/tensorflow/lite/tools/benchmark/README.md b/tensorflow/lite/tools/benchmark/README.md
index a71a2fa..a4d9c87 100644
--- a/tensorflow/lite/tools/benchmark/README.md
+++ b/tensorflow/lite/tools/benchmark/README.md
@@ -11,6 +11,11 @@
for iOS please use the
[iOS benchmark app](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/benchmark/ios).
+An experimental Android APK wrapper for the benchmark model utility offers more
+faithful execution behavior on Android (via a foreground Activity). It is
+located
+[here](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/benchmark/android).
+
## Parameters
The binary takes the following required parameters:
diff --git a/tensorflow/lite/tools/benchmark/android/AndroidManifest.xml b/tensorflow/lite/tools/benchmark/android/AndroidManifest.xml
new file mode 100644
index 0000000..7cdca28
--- /dev/null
+++ b/tensorflow/lite/tools/benchmark/android/AndroidManifest.xml
@@ -0,0 +1,42 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+ 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+ package="org.tensorflow.lite.benchmark">
+
+ <!-- Necessary for loading custom models from disk. -->
+ <uses-permission android:name="android.permission.READ_EXTERNAL_STORAGE"/>
+
+ <!-- Target SDK 21 (<23) to avoid the need for requesting storage
+ permissions. This APK will almost always be used from the command-line
+ anyway, and be expicitly installed by the developer. -->
+ <uses-sdk
+ android:minSdkVersion="21"
+ android:targetSdkVersion="21" />
+
+ <application>
+ <!-- This Activity runs the TensorFlow Lite benchmark at creation, using
+ a provided set of arguments, then immediately terminates. -->
+ <activity android:name="org.tensorflow.lite.benchmark.BenchmarkModelActivity"
+ android:screenOrientation="portrait"
+ android:label="TFLite Benchmark"
+ android:theme="@android:style/Theme.NoDisplay"
+ android:exported="true"
+ android:noHistory="true" />
+ </application>
+
+</manifest>
diff --git a/tensorflow/lite/tools/benchmark/android/BUILD b/tensorflow/lite/tools/benchmark/android/BUILD
new file mode 100644
index 0000000..a291eff
--- /dev/null
+++ b/tensorflow/lite/tools/benchmark/android/BUILD
@@ -0,0 +1,44 @@
+# Description:
+# BenchmarkModel Android harness for TensorFlow Lite benchmarks.
+package(default_visibility = ["//visibility:private"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(["LICENSE"])
+
+load("//tensorflow/lite:build_def.bzl", "tflite_jni_binary")
+load("@build_bazel_rules_android//android:rules.bzl", "android_binary")
+
+# See README.md for details about building and executing this benchmark.
+android_binary(
+ name = "benchmark_model",
+ srcs = glob([
+ "src/**/*.java",
+ ]),
+ custom_package = "org.tensorflow.lite.benchmark",
+ manifest = "AndroidManifest.xml",
+ # In some platforms we don't have an Android SDK/NDK and this target
+ # can't be built. We need to prevent the build system from trying to
+ # use the target in that case.
+ tags = ["manual"],
+ deps = [":tensorflowlite_benchmark_native"],
+)
+
+tflite_jni_binary(
+ name = "libtensorflowlite_benchmark.so",
+ srcs = glob([
+ "jni/**/*.cc",
+ "jni/**/*.h",
+ ]),
+ deps = [
+ "//tensorflow/lite/java/jni",
+ "//tensorflow/lite/tools/benchmark:benchmark_tflite_model_lib",
+ "//tensorflow/lite/tools/benchmark:logging",
+ ],
+)
+
+cc_library(
+ name = "tensorflowlite_benchmark_native",
+ srcs = ["libtensorflowlite_benchmark.so"],
+ visibility = ["//visibility:private"],
+)
diff --git a/tensorflow/lite/tools/benchmark/android/README.md b/tensorflow/lite/tools/benchmark/android/README.md
new file mode 100644
index 0000000..f5b67e3
--- /dev/null
+++ b/tensorflow/lite/tools/benchmark/android/README.md
@@ -0,0 +1,65 @@
+# TFLite Android Model Benchmark Tool
+
+## Description
+
+This Android benchmark app is a simple wrapper around the TensorFlow Lite
+[command-line benchmark utility](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/benchmark).
+
+Pushing and executing binaries directly on Android is a valid approach to
+benchmarking, but it can result in subtle (but observable) differences in
+performance relative to execution within an actual Android app. In particular,
+Android's scheduler tailors behavior based on thread and process priorities,
+which differ between a foreground Activity/Application and a regular background
+binary executed via `adb shell ...`. This tailored behavior is most evident when
+enabling multi-threaded CPU execution with TensorFlow Lite.
+
+To that end, this app offers perhaps a more faithful view of runtime performance
+that developers can expected when deploying TensorFlow Lite with their
+application.
+
+## To build/install/run
+
+(0) Refer to
+https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android
+to edit the `WORKSPACE` to configure the android NDK/SDK.
+
+(1) Build for your specific platform, e.g.:
+
+```
+bazel build -c opt \
+ --config=android_arm64 \
+ --cxxopt='--std=c++11' \
+ tensorflow/lite/tools/benchmark/android:benchmark_model
+```
+
+(2) Connect your phone. Install the benchmark APK to your phone with adb:
+
+```
+adb install -r -d bazel-bin/tensorflow/lite/tools/benchmark/android/benchmark_model.apk
+```
+
+(3) Push the compute graph that you need to test.
+
+```
+adb push mobilenet_quant_v1_224.tflite /data/local/tmp
+```
+
+(4) Run the benchmark. Additional command-line flags are documented
+[here](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/benchmark/README.md)
+and can be appended to the `args` string alongside the required `--graph` flag
+(note that all args must be nested in the single quoted string that follows the
+args key).
+
+```
+adb shell am start -S -n
+ org.tensorflow.lite.benchmark/org.tensorflow.lite.benchmark.BenchmarkModelActivity \
+ --es args '"--graph=/data/local/tmp/mobilenet_quant_v1_224.tflite --num_threads=4"'
+```
+
+(5) The results will be available in Android's logcat, e.g.:
+
+```
+adb logcat | grep "Average inference"
+
+... tflite : Average inference timings in us: Warmup: 91471, Init: 4108, Inference: 80660.1
+```
diff --git a/tensorflow/lite/tools/benchmark/android/jni/benchmark_model_jni.cc b/tensorflow/lite/tools/benchmark/android/jni/benchmark_model_jni.cc
new file mode 100644
index 0000000..ee67bda
--- /dev/null
+++ b/tensorflow/lite/tools/benchmark/android/jni/benchmark_model_jni.cc
@@ -0,0 +1,92 @@
+/* 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.
+==============================================================================*/
+
+#include <jni.h>
+
+#include <sstream>
+#include <string>
+
+#include "tensorflow/lite/tools/benchmark/benchmark_tflite_model.h"
+#include "tensorflow/lite/tools/benchmark/logging.h"
+
+#ifdef __ANDROID__
+#include <android/log.h>
+#endif
+
+namespace tflite {
+namespace benchmark {
+namespace {
+
+class AndroidBenchmarkLoggingListener : public BenchmarkListener {
+ void OnBenchmarkEnd(const BenchmarkResults& results) override {
+ auto inference_us = results.inference_time_us();
+ auto init_us = results.startup_latency_us();
+ auto warmup_us = results.warmup_time_us();
+ std::stringstream results_output;
+ results_output << "Average inference timings in us: "
+ << "Warmup: " << warmup_us.avg() << ", "
+ << "Init: " << init_us << ", "
+ << "Inference: " << inference_us.avg();
+#ifdef __ANDROID__
+ __android_log_print(ANDROID_LOG_ERROR, "tflite", "%s",
+ results_output.str().c_str());
+#else
+ fprintf(stderr, "%s", results_output.str().c_str());
+#endif
+ }
+};
+
+void Run(int argc, char** argv) {
+ BenchmarkTfLiteModel benchmark;
+ AndroidBenchmarkLoggingListener listener;
+ benchmark.AddListener(&listener);
+ benchmark.Run(argc, argv);
+}
+
+} // namespace
+} // namespace benchmark
+} // namespace tflite
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_benchmark_BenchmarkModel_nativeRun(JNIEnv* env,
+ jclass clazz,
+ jstring args_obj) {
+ const char* args_chars = env->GetStringUTFChars(args_obj, nullptr);
+
+ // Split the args string into individual arg tokens.
+ std::istringstream iss(args_chars);
+ std::vector<std::string> args_split{std::istream_iterator<std::string>(iss),
+ {}};
+
+ // Construct a fake argv command-line object for the benchmark.
+ std::vector<char*> argv;
+ std::string arg0 = "(BenchmarkModelAndroid)";
+ argv.push_back(const_cast<char*>(arg0.data()));
+ for (auto& arg : args_split) {
+ argv.push_back(const_cast<char*>(arg.data()));
+ }
+
+ tflite::benchmark::Run(static_cast<int>(argv.size()), argv.data());
+
+ env->ReleaseStringUTFChars(args_obj, args_chars);
+}
+
+#ifdef __cplusplus
+} // extern "C"
+#endif // __cplusplus
diff --git a/tensorflow/lite/tools/benchmark/android/src/org/tensorflow/lite/benchmark/BenchmarkModel.java b/tensorflow/lite/tools/benchmark/android/src/org/tensorflow/lite/benchmark/BenchmarkModel.java
new file mode 100644
index 0000000..a6cf8d7
--- /dev/null
+++ b/tensorflow/lite/tools/benchmark/android/src/org/tensorflow/lite/benchmark/BenchmarkModel.java
@@ -0,0 +1,32 @@
+/* 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.
+==============================================================================*/
+
+package org.tensorflow.lite.benchmark;
+
+/** Helper class for running a native TensorFlow Lite benchmark. */
+class BenchmarkModel {
+ static {
+ System.loadLibrary("tensorflowlite_benchmark");
+ }
+
+ // Executes a standard TensorFlow Lite benchmark according to the provided args.
+ //
+ // Note that {@code args} will be split by the native execution code.
+ public static void run(String args) {
+ nativeRun(args);
+ }
+
+ private static native void nativeRun(String args);
+}
diff --git a/tensorflow/lite/tools/benchmark/android/src/org/tensorflow/lite/benchmark/BenchmarkModelActivity.java b/tensorflow/lite/tools/benchmark/android/src/org/tensorflow/lite/benchmark/BenchmarkModelActivity.java
new file mode 100644
index 0000000..12410ad
--- /dev/null
+++ b/tensorflow/lite/tools/benchmark/android/src/org/tensorflow/lite/benchmark/BenchmarkModelActivity.java
@@ -0,0 +1,44 @@
+/* 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.
+==============================================================================*/
+
+package org.tensorflow.lite.benchmark;
+
+import android.app.Activity;
+import android.content.Intent;
+import android.os.Bundle;
+import android.util.Log;
+
+/** Main {@code Activity} class for the benchmark app. */
+public class BenchmarkModelActivity extends Activity {
+
+ private static final String TAG = "tflite_BenchmarkModelActivity";
+
+ private static final String ARGS_INTENT_KEY_0 = "args";
+ private static final String ARGS_INTENT_KEY_1 = "--args";
+
+ @Override
+ public void onCreate(Bundle savedInstanceState) {
+ super.onCreate(savedInstanceState);
+
+ Intent intent = getIntent();
+ Bundle bundle = intent.getExtras();
+ String args = bundle.getString(ARGS_INTENT_KEY_0, bundle.getString(ARGS_INTENT_KEY_1));
+ Log.i(TAG, "Running TensorFlow Lite benchmark with args: " + args);
+
+ BenchmarkModel.run(args);
+
+ finish();
+ }
+}
diff --git a/tensorflow/lite/tools/benchmark/benchmark_test.cc b/tensorflow/lite/tools/benchmark/benchmark_test.cc
index 8191fbc..a4f8301 100644
--- a/tensorflow/lite/tools/benchmark/benchmark_test.cc
+++ b/tensorflow/lite/tools/benchmark/benchmark_test.cc
@@ -18,6 +18,7 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>
+#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/testing/util.h"
#include "tensorflow/lite/tools/benchmark/benchmark_tflite_model.h"
#include "tensorflow/lite/tools/benchmark/command_line_flags.h"
@@ -47,6 +48,15 @@
return params;
}
+class TestBenchmark : public BenchmarkTfLiteModel {
+ public:
+ explicit TestBenchmark(BenchmarkParams params)
+ : BenchmarkTfLiteModel(std::move(params)) {}
+ const tflite::Interpreter* GetInterpreter() { return interpreter.get(); }
+
+ void Prepare() { PrepareInputsAndOutputs(); }
+};
+
TEST(BenchmarkTest, DoesntCrash) {
ASSERT_THAT(g_model_path, testing::NotNull());
@@ -54,6 +64,37 @@
benchmark.Run();
}
+TEST(BenchmarkTest, ParametersArePopulatedWhenInputShapeIsNotSpecified) {
+ ASSERT_THAT(g_model_path, testing::NotNull());
+
+ TestBenchmark benchmark(CreateParams());
+ benchmark.Init();
+ benchmark.Prepare();
+
+ auto interpreter = benchmark.GetInterpreter();
+ auto inputs = interpreter->inputs();
+ ASSERT_GE(inputs.size(), 1);
+ auto input_tensor = interpreter->tensor(inputs[0]);
+
+ std::vector<char> input_bytes;
+ input_bytes.reserve(input_tensor->bytes);
+ for (size_t i = 0; i < input_tensor->bytes; i++) {
+ input_bytes.push_back(input_tensor->data.raw_const[i]);
+ }
+ benchmark.Prepare();
+
+ // Expect data is not the same.
+ EXPECT_EQ(input_bytes.size(), input_tensor->bytes);
+ bool is_same = true;
+ for (size_t i = 0; i < input_tensor->bytes; i++) {
+ if (input_bytes[i] != input_tensor->data.raw_const[i]) {
+ is_same = false;
+ break;
+ }
+ }
+ EXPECT_FALSE(is_same);
+}
+
} // namespace
} // namespace benchmark
} // namespace tflite
diff --git a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
index 16f7087..32cf4e4 100644
--- a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
+++ b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc
@@ -30,7 +30,7 @@
#include "tensorflow/lite/tools/benchmark/logging.h"
#ifdef GEMMLOWP_PROFILING
-#include "third_party/gemmlowp/profiling/profiler.h"
+#include "gemmlowp/profiling/profiler.h"
#endif
#ifdef TFLITE_CUSTOM_OPS_HEADER
@@ -181,6 +181,15 @@
return true;
}
+std::vector<int> TfLiteIntArrayToVector(const TfLiteIntArray* int_array) {
+ std::vector<int> values;
+ values.reserve(int_array->size);
+ for (size_t i = 0; i < int_array->size; i++) {
+ values.push_back(int_array->data[i]);
+ }
+ return values;
+}
+
} // namespace
BenchmarkParams BenchmarkTfLiteModel::DefaultParams() {
@@ -250,12 +259,10 @@
void BenchmarkTfLiteModel::PrepareInputsAndOutputs() {
auto interpreter_inputs = interpreter->inputs();
// Set the values of the input tensors.
- for (int j = 0; j < inputs.size(); ++j) {
- const InputLayerInfo& input = inputs[j];
+ for (int j = 0; j < interpreter_inputs.size(); ++j) {
int i = interpreter_inputs[j];
TfLiteTensor* t = interpreter->tensor(i);
- std::vector<int> sizes = input.shape;
-
+ std::vector<int> sizes = TfLiteIntArrayToVector(t->dims);
// TODO(ahentz): below we ignore the O-th dimension (number of batches).
if (t->type == kTfLiteFloat32) {
FillRandomValue<float>(
@@ -274,6 +281,11 @@
interpreter->typed_tensor<uint8_t>(i),
std::vector<int>(sizes.begin() + 1, sizes.end()),
[]() { return static_cast<uint8_t>(rand()) % 255; });
+ } else if (t->type == kTfLiteInt8) {
+ FillRandomValue<int8_t>(
+ interpreter->typed_tensor<int8_t>(i),
+ std::vector<int>(sizes.begin() + 1, sizes.end()),
+ []() { return static_cast<int8_t>(rand()) % 255 - 127; });
} else if (t->type == kTfLiteString) {
tflite::DynamicBuffer buffer;
FillRandomString(&buffer, sizes, []() {
diff --git a/tensorflow/lite/tools/benchmark/ios/README.md b/tensorflow/lite/tools/benchmark/ios/README.md
index 3dc29d9..fed9e7e 100644
--- a/tensorflow/lite/tools/benchmark/ios/README.md
+++ b/tensorflow/lite/tools/benchmark/ios/README.md
@@ -41,3 +41,14 @@
- Now try running the app. The app has a single button that runs the benchmark
on the model and displays results in a text view below.
+
+## Profiling
+
+If you want detailed profiling, use the following command:
+
+```bash
+tensorflow/lite/build_ios_universal_lib.sh -p
+```
+
+Then following the same steps above and run the benchmark app. You will see the
+detailed profiling results in the outputs.
diff --git a/tensorflow/lite/tools/make/Makefile b/tensorflow/lite/tools/make/Makefile
index 363a069..994f660 100644
--- a/tensorflow/lite/tools/make/Makefile
+++ b/tensorflow/lite/tools/make/Makefile
@@ -52,6 +52,7 @@
# generate things like the protobuf compiler that require that), so all of
# these settings are for the target compiler.
CXXFLAGS := -O3 -DNDEBUG
+CXXFLAGS += $(EXTRA_CXXFLAGS)
CCFLAGS := ${CXXFLAGS}
CXXFLAGS += --std=c++11
CFLAGS :=
@@ -116,8 +117,8 @@
tensorflow/lite/nnapi_delegate.cc
else
CORE_CC_EXCLUDE_SRCS += \
-tensorflow/contrib/lite/mmap_allocation_disabled.cc \
-tensorflow/contrib/lite/nnapi_delegate_disabled.cc
+tensorflow/lite/mmap_allocation_disabled.cc \
+tensorflow/lite/nnapi_delegate_disabled.cc
endif
# Filter out all the excluded files.
TF_LITE_CC_SRCS := $(filter-out $(CORE_CC_EXCLUDE_SRCS), $(CORE_CC_ALL_SRCS))
diff --git a/tensorflow/lite/tools/make/build_ios_universal_lib.sh b/tensorflow/lite/tools/make/build_ios_universal_lib.sh
index 6e0d262..8b617ef 100755
--- a/tensorflow/lite/tools/make/build_ios_universal_lib.sh
+++ b/tensorflow/lite/tools/make/build_ios_universal_lib.sh
@@ -19,20 +19,36 @@
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR/../../../.."
+usage() {
+ echo "Usage: $(basename "$0") [-a]"
+ echo "-a [build_arch] build for specified arch comma separate for multiple archs (eg: x86_64 arm64)"
+ echo " default is [x86_64 armv7 armv7s arm64]"
+ echo "-p enable profiling"
+ exit 1
+}
+
+profiling_args=""
+BUILD_ARCHS="x86_64 armv7 armv7s arm64"
+while getopts "a:p" opt_name; do
+ case "$opt_name" in
+ a) BUILD_ARCHS="${OPTARG}";;
+ p) profiling_args='-DGEMMLOWP_PROFILING,-DTFLITE_PROFILING_ENABLED';;
+ *) usage;;
+ esac
+done
+shift $(($OPTIND - 1))
+
# Build library for supported architectures and packs them in a fat binary.
make_library() {
- for arch in x86_64 armv7 armv7s arm64
+ LIBS=""
+ for arch in $BUILD_ARCHS
do
make -f tensorflow/lite/tools/make/Makefile TARGET=ios TARGET_ARCH=${arch} \
- -j 8
+ EXTRA_CXXFLAGS=$profiling_args -j 8
+ LIBS="${LIBS} tensorflow/lite/tools/make/gen/ios_${arch}/lib/${1}"
done
mkdir -p tensorflow/lite/tools/make/gen/lib
- lipo \
- tensorflow/lite/tools/make/gen/ios_x86_64/lib/${1} \
- tensorflow/lite/tools/make/gen/ios_armv7/lib/${1} \
- tensorflow/lite/tools/make/gen/ios_armv7s/lib/${1} \
- tensorflow/lite/tools/make/gen/ios_arm64/lib/${1} \
- -create \
+ lipo $LIBS -create \
-output tensorflow/lite/tools/make/gen/lib/${1}
}
diff --git a/tensorflow/lite/tutorials/post_training_quant.ipynb b/tensorflow/lite/tutorials/post_training_quant.ipynb
index 3ff145d..394ab07 100644
--- a/tensorflow/lite/tutorials/post_training_quant.ipynb
+++ b/tensorflow/lite/tutorials/post_training_quant.ipynb
@@ -235,9 +235,9 @@
"id": "AT8BgkKmljOy"
},
"source": [
- "Using the python `TocoConverter`, the saved model can be converted into a TFLite model.\n",
+ "Using the python `TFLiteConverter`, the saved model can be converted into a TFLite model.\n",
"\n",
- "First load the model using the `TocoConverter`:"
+ "First load the model using the `TFLiteConverter`:"
]
},
{
@@ -252,7 +252,7 @@
"source": [
"import tensorflow as tf\n",
"tf.enable_eager_execution()\n",
- "converter = tf.lite.TocoConverter.from_saved_model(saved_model_dir)\n",
+ "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)\n",
"tflite_model = converter.convert()"
]
},
@@ -648,7 +648,7 @@
"graph_def_file = pathlib.Path(archive_path).parent/\"resnet_v2_101_299_frozen.pb\"\n",
"input_arrays = [\"input\"] \n",
"output_arrays = [\"output\"]\n",
- "converter = tf.lite.TocoConverter.from_frozen_graph(\n",
+ "converter = tf.lite.TFLiteConverter.from_frozen_graph(\n",
" str(graph_def_file), input_arrays, output_arrays, input_shapes={\"input\":[1,299,299,3]})\n",
"converter.post_training_quantize = True\n",
"resnet_tflite_file = graph_def_file.parent/\"resnet_v2_101_quantized.tflite\"\n",
diff --git a/tensorflow/opensource_only.files b/tensorflow/opensource_only.files
index 8ff1645..0af84f8 100644
--- a/tensorflow/opensource_only.files
+++ b/tensorflow/opensource_only.files
@@ -1,3 +1,42 @@
+tensorflow/contrib/tpu/profiler/pip_package/BUILD
+tensorflow/contrib/tpu/profiler/pip_package/setup.py
+tensorflow/contrib/tpu/profiler/pip_package/README
+tensorflow/contrib/tpu/profiler/pip_package/build_pip_package.sh
+tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py
+tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/__init__.py
+tensorflow/contrib/mpi/BUILD
+tensorflow/tools/ci_build/remote/BUILD
+tensorflow/tools/pip_package/README
+tensorflow/tools/pip_package/MANIFEST.in
+tensorflow/tools/pip_package/simple_console.py
+tensorflow/tools/pip_package/build_pip_package.sh
+tensorflow/tools/pip_package/check_load_py_test.py
+tensorflow/tools/pip_package/pip_smoke_test.py
+tensorflow/tools/pip_package/simple_console_for_windows.py
+tensorflow/tools/pip_package/setup.py
+tensorflow/tools/pip_package/BUILD
+tensorflow/tools/lib_package/concat_licenses.sh
+tensorflow/tools/lib_package/libtensorflow_test.c
+tensorflow/tools/lib_package/LibTensorFlowTest.java
+tensorflow/tools/lib_package/BUILD
+tensorflow/tools/lib_package/libtensorflow_test.sh
+tensorflow/tools/lib_package/README.md
+tensorflow/tools/lib_package/libtensorflow_java_test.sh
+tensorflow/tools/def_file_filter/def_file_filter_configure.bzl
+tensorflow/tools/def_file_filter/BUILD
+tensorflow/tools/def_file_filter/BUILD.tpl
+tensorflow/tools/def_file_filter/def_file_filter.py.tpl
+tensorflow/third_party/mkl/MKL_LICENSE
+tensorflow/third_party/mkl/LICENSE
+tensorflow/third_party/mkl/BUILD
+tensorflow/third_party/mkl/mkl.BUILD
+tensorflow/third_party/mkl/build_defs.bzl
+tensorflow/third_party/backports_weakref.BUILD
+tensorflow/third_party/toolchains/clang6/BUILD
+tensorflow/third_party/toolchains/clang6/README.md
+tensorflow/third_party/toolchains/clang6/repo.bzl
+tensorflow/third_party/toolchains/clang6/CROSSTOOL.tpl
+tensorflow/third_party/toolchains/clang6/clang.BUILD
tensorflow/third_party/toolchains/preconfig/ubuntu14.04/py3/BUILD
tensorflow/third_party/toolchains/preconfig/ubuntu14.04/gcc-nvcc-cuda9.0/BUILD
tensorflow/third_party/toolchains/preconfig/ubuntu14.04/gcc-nvcc/BUILD
@@ -10,8 +49,195 @@
tensorflow/third_party/toolchains/preconfig/generate/workspace.bzl
tensorflow/third_party/toolchains/preconfig/generate/containers.bzl
tensorflow/third_party/toolchains/preconfig/generate/generate.bzl
+tensorflow/third_party/toolchains/preconfig/generate/archives.bzl
tensorflow/third_party/toolchains/preconfig/generate/BUILD
tensorflow/third_party/toolchains/preconfig/win_1803/bazel_018/BUILD
tensorflow/third_party/toolchains/preconfig/win_1803/bazel_018/dummy_toolchain.bzl
tensorflow/third_party/toolchains/preconfig/win_1803/py36/BUILD
-tensorflow/third_party/toolchains/preconfig/win_1803/BUILD
\ No newline at end of file
+tensorflow/third_party/toolchains/preconfig/win_1803/BUILD
+tensorflow/third_party/toolchains/gpus/cuda/build_defs.bzl
+tensorflow/third_party/toolchains/gpus/cuda/BUILD
+tensorflow/third_party/toolchains/gpus/cuda/cuda/cuda_config.h
+tensorflow/third_party/toolchains/gpus/crosstool/BUILD
+tensorflow/third_party/toolchains/gpus/crosstool/CROSSTOOL
+tensorflow/third_party/toolchains/gpus/py/BUILD
+tensorflow/third_party/toolchains/cpus/arm/arm_compiler_configure.bzl
+tensorflow/third_party/toolchains/cpus/arm/CROSSTOOL.tpl
+tensorflow/third_party/toolchains/cpus/arm/BUILD
+tensorflow/third_party/toolchains/cpus/py3/BUILD
+tensorflow/third_party/toolchains/cpus/py/BUILD
+tensorflow/third_party/toolchains/BUILD
+tensorflow/third_party/nccl/remote.BUILD.tpl
+tensorflow/third_party/nccl/archive.BUILD
+tensorflow/third_party/nccl/LICENSE
+tensorflow/third_party/nccl/system.BUILD.tpl
+tensorflow/third_party/nccl/nccl_configure.bzl
+tensorflow/third_party/nccl/build_defs.bzl.tpl
+tensorflow/third_party/nccl/BUILD
+tensorflow/third_party/gpus/BUILD
+tensorflow/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_rocm.tpl
+tensorflow/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc.tpl
+tensorflow/third_party/gpus/crosstool/CROSSTOOL.tpl
+tensorflow/third_party/gpus/crosstool/CROSSTOOL_hipcc.tpl
+tensorflow/third_party/gpus/crosstool/LICENSE
+tensorflow/third_party/gpus/crosstool/remote.BUILD.tpl
+tensorflow/third_party/gpus/crosstool/windows/msvc_wrapper_for_nvcc.py.tpl
+tensorflow/third_party/gpus/crosstool/BUILD.tpl
+tensorflow/third_party/gpus/crosstool/BUILD
+tensorflow/third_party/gpus/cuda/LICENSE
+tensorflow/third_party/gpus/cuda/BUILD.tpl
+tensorflow/third_party/gpus/cuda/BUILD.windows.tpl
+tensorflow/third_party/gpus/cuda/cuda_config.h.tpl
+tensorflow/third_party/gpus/cuda/remote.BUILD.tpl
+tensorflow/third_party/gpus/cuda/BUILD
+tensorflow/third_party/gpus/cuda/build_defs.bzl.tpl
+tensorflow/third_party/gpus/rocm/rocm_config.h.tpl
+tensorflow/third_party/gpus/rocm/BUILD
+tensorflow/third_party/gpus/rocm/BUILD.tpl
+tensorflow/third_party/gpus/rocm/build_defs.bzl.tpl
+tensorflow/third_party/gpus/cuda_configure.bzl
+tensorflow/third_party/gpus/rocm_configure.bzl
+tensorflow/third_party/snappy.BUILD
+tensorflow/third_party/cython.BUILD
+tensorflow/third_party/farmhash.BUILD
+tensorflow/third_party/eigen3/Eigen/Cholesky
+tensorflow/third_party/eigen3/Eigen/QR
+tensorflow/third_party/eigen3/Eigen/LU
+tensorflow/third_party/eigen3/Eigen/Core
+tensorflow/third_party/eigen3/Eigen/SVD
+tensorflow/third_party/eigen3/Eigen/Eigenvalues
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/FixedPoint
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/Tensor
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/PacketMathAVX512.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/TypeCastingAVX512.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/PacketMathAVX2.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/MatMatProduct.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/FixedPointTypes.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/MatVecProduct.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/MatMatProductNEON.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/TypeCastingAVX2.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/src/FixedPoint/MatMatProductAVX2.h
+tensorflow/third_party/eigen3/unsupported/Eigen/CXX11/ThreadPool
+tensorflow/third_party/eigen3/unsupported/Eigen/SpecialFunctions
+tensorflow/third_party/eigen3/unsupported/Eigen/MatrixFunctions
+tensorflow/third_party/eigen3/LICENSE
+tensorflow/third_party/eigen3/BUILD
+tensorflow/third_party/systemlibs/build_defs.bzl.tpl
+tensorflow/third_party/systemlibs/absl_py.BUILD
+tensorflow/third_party/systemlibs/curl.BUILD
+tensorflow/third_party/systemlibs/termcolor.BUILD
+tensorflow/third_party/systemlibs/absl_py.absl.flags.BUILD
+tensorflow/third_party/systemlibs/grpc.BUILD
+tensorflow/third_party/systemlibs/swig.BUILD
+tensorflow/third_party/systemlibs/protobuf.bzl
+tensorflow/third_party/systemlibs/protobuf.BUILD
+tensorflow/third_party/systemlibs/BUILD
+tensorflow/third_party/systemlibs/google_cloud_cpp.BUILD
+tensorflow/third_party/systemlibs/astor.BUILD
+tensorflow/third_party/systemlibs/six.BUILD
+tensorflow/third_party/systemlibs/absl_py.absl.testing.BUILD
+tensorflow/third_party/systemlibs/boringssl.BUILD
+tensorflow/third_party/systemlibs/nsync.BUILD
+tensorflow/third_party/systemlibs/google_cloud_cpp.google.cloud.bigtable.BUILD
+tensorflow/third_party/systemlibs/gif.BUILD
+tensorflow/third_party/systemlibs/pcre.BUILD
+tensorflow/third_party/systemlibs/BUILD.tpl
+tensorflow/third_party/systemlibs/snappy.BUILD
+tensorflow/third_party/systemlibs/gast.BUILD
+tensorflow/third_party/systemlibs/cython.BUILD
+tensorflow/third_party/systemlibs/double_conversion.BUILD
+tensorflow/third_party/systemlibs/zlib.BUILD
+tensorflow/third_party/systemlibs/jsoncpp.BUILD
+tensorflow/third_party/systemlibs/re2.BUILD
+tensorflow/third_party/systemlibs/lmdb.BUILD
+tensorflow/third_party/systemlibs/googleapis.BUILD
+tensorflow/third_party/systemlibs/png.BUILD
+tensorflow/third_party/systemlibs/syslibs_configure.bzl
+tensorflow/third_party/systemlibs/sqlite.BUILD
+tensorflow/third_party/python_runtime/BUILD
+tensorflow/third_party/sycl/crosstool/BUILD
+tensorflow/third_party/ngraph/LICENSE
+tensorflow/third_party/ngraph/tbb.BUILD
+tensorflow/third_party/ngraph/BUILD
+tensorflow/third_party/ngraph/ngraph.BUILD
+tensorflow/third_party/ngraph/build_defs.bzl
+tensorflow/third_party/ngraph/NGRAPH_LICENSE
+tensorflow/third_party/ngraph/ngraph_tf.BUILD
+tensorflow/third_party/ngraph/nlohmann_json.BUILD
+tensorflow/third_party/clang_toolchain/download_clang.bzl
+tensorflow/third_party/clang_toolchain/BUILD
+tensorflow/third_party/clang_toolchain/cc_configure_clang.bzl
+tensorflow/third_party/gast.BUILD
+tensorflow/third_party/llvm/BUILD
+tensorflow/third_party/llvm/expand_cmake_vars.py
+tensorflow/third_party/llvm/llvm.autogenerated.BUILD
+tensorflow/third_party/llvm/llvm.bzl
+tensorflow/third_party/icu/udata.patch
+tensorflow/third_party/fft2d/BUILD
+tensorflow/third_party/fft2d/fft.h
+tensorflow/third_party/fft2d/LICENSE
+tensorflow/third_party/fft2d/fft2d.BUILD
+tensorflow/third_party/boringssl/BUILD
+tensorflow/third_party/mpi/.gitignore
+tensorflow/third_party/mpi/BUILD
+tensorflow/third_party/tensorrt/LICENSE
+tensorflow/third_party/tensorrt/BUILD
+tensorflow/third_party/tensorrt/build_defs.bzl.tpl
+tensorflow/third_party/tensorrt/BUILD.tpl
+tensorflow/third_party/tensorrt/tensorrt_configure.bzl
+tensorflow/third_party/kafka/config.patch
+tensorflow/third_party/kafka/BUILD
+tensorflow/third_party/android/BUILD
+tensorflow/third_party/android/android.bzl.tpl
+tensorflow/third_party/android/android_configure.bzl
+tensorflow/third_party/android/android_configure.BUILD.tpl
+tensorflow/third_party/tflite_smartreply.BUILD
+tensorflow/third_party/mkl_dnn/LICENSE
+tensorflow/third_party/mkl_dnn/mkldnn.BUILD
+tensorflow/third_party/pcre.BUILD
+tensorflow/third_party/linenoise.BUILD
+tensorflow/third_party/sqlite.BUILD
+tensorflow/third_party/common.bzl
+tensorflow/third_party/com_google_absl.BUILD
+tensorflow/third_party/pprof.BUILD
+tensorflow/third_party/BUILD
+tensorflow/third_party/tflite_mobilenet_quant.BUILD
+tensorflow/third_party/lmdb.BUILD
+tensorflow/third_party/git/BUILD.tpl
+tensorflow/third_party/git/BUILD
+tensorflow/third_party/git/git_configure.bzl
+tensorflow/third_party/protobuf/BUILD
+tensorflow/third_party/tflite_mobilenet.BUILD
+tensorflow/third_party/py/BUILD
+tensorflow/third_party/py/BUILD.tpl
+tensorflow/third_party/py/remote.BUILD.tpl
+tensorflow/third_party/py/numpy/BUILD
+tensorflow/third_party/py/python_configure.bzl
+tensorflow/third_party/termcolor.BUILD
+tensorflow/third_party/png_fix_rpi.patch
+tensorflow/third_party/swig.BUILD
+tensorflow/third_party/astor.BUILD
+tensorflow/third_party/grpc/BUILD
+tensorflow/third_party/curl.BUILD
+tensorflow/third_party/arm_neon_2_x86_sse.BUILD
+tensorflow/third_party/png.BUILD
+tensorflow/third_party/googleapis.BUILD
+tensorflow/third_party/mpi_collectives/BUILD
+tensorflow/third_party/nanopb.BUILD
+tensorflow/third_party/gif.BUILD
+tensorflow/third_party/double_conversion.BUILD
+tensorflow/third_party/six.BUILD
+tensorflow/third_party/tflite_mobilenet_float.BUILD
+tensorflow/third_party/repo.bzl
+tensorflow/third_party/codegen.BUILD
+tensorflow/third_party/cub.BUILD
+tensorflow/third_party/jsoncpp.BUILD
+tensorflow/third_party/tflite_ovic_testdata.BUILD
+tensorflow/third_party/libxsmm.BUILD
+tensorflow/third_party/zlib.BUILD
+tensorflow/third_party/eigen.BUILD
+tensorflow/stream_executor/BUILD
+tensorflow/api_template_v1.__init__.py
+tensorflow/compat_template_v1.__init__.py
+tensorflow/api_template.__init__.py
+tensorflow/__init__.py
\ No newline at end of file
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 1384a2e..8c0e95e 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -6,6 +6,7 @@
visibility = [
"//engedu/ml/tf_from_scratch:__pkg__",
+ "//third_party/cloud_tpu/convergence_tools:__subpackages__",
"//tensorflow:internal",
"//tensorflow/lite/toco/python:__pkg__",
"//tensorflow_models:__subpackages__",
@@ -145,6 +146,7 @@
"//tensorflow/lite/python:lite",
"//tensorflow/python/compat",
"//tensorflow/python/data",
+ "//tensorflow/python/distribute",
"//tensorflow/python/distribute:estimator_training",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/feature_column:feature_column_py",
@@ -3031,6 +3033,7 @@
"//tensorflow/python/eager:wrap_function",
"//tensorflow/python/ops/distributions",
"//tensorflow/python/ops/linalg",
+ "//tensorflow/python/ops/ragged",
],
)
@@ -3235,6 +3238,7 @@
srcs = ["ops/control_flow_ops_test.py"],
additional_deps = [
":array_ops",
+ ":cond_v2",
":control_flow_ops",
":embedding_ops",
":framework_for_generated_wrappers",
@@ -3250,6 +3254,7 @@
":util",
":variable_scope",
":variables",
+ ":while_v2",
"//tensorflow/python/eager:def_function",
],
)
@@ -3512,13 +3517,13 @@
exclude = [
"**/*test*",
"training/checkpointable/**/*.py",
+ "training/saving/**/*.py",
# The following targets have their own build rules (same name as the
# file):
"training/basic_session_run_hooks.py",
"training/checkpoint_management.py",
"training/distribute.py",
"training/distribution_strategy_context.py",
- "training/saveable_object.py",
"training/saver.py",
"training/session_run_hook.py",
"training/training_util.py",
@@ -3594,12 +3599,6 @@
)
py_library(
- name = "saveable_object",
- srcs = ["training/saveable_object.py"],
- srcs_version = "PY2AND3",
-)
-
-py_library(
name = "checkpoint_management",
srcs = ["training/checkpoint_management.py"],
deps = [
@@ -3652,7 +3651,6 @@
":platform",
":pywrap_tensorflow",
":resource_variable_ops",
- ":saveable_object",
":session",
":state_ops",
":string_ops",
@@ -3662,6 +3660,8 @@
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/training/checkpointable:base",
+ "//tensorflow/python/training/saving:saveable_object",
+ "//tensorflow/python/training/saving:saveable_object_util",
"//third_party/py/numpy",
"@six_archive//:six",
],
@@ -3835,6 +3835,7 @@
srcs_version = "PY2AND3",
deps = [
":util",
+ "//tensorflow/python:framework_ops",
"//tensorflow/python/eager:context",
"@six_archive//:six",
],
@@ -5321,6 +5322,20 @@
)
py_test(
+ name = "quantized_ops_test",
+ size = "small",
+ srcs = ["ops/quantized_ops_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_windows"],
+ deps = [
+ ":array_ops",
+ ":client_testlib",
+ ":framework_for_generated_wrappers",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
name = "quantized_conv_ops_test",
size = "small",
srcs = ["ops/quantized_conv_ops_test.py"],
diff --git a/tensorflow/python/__init__.py b/tensorflow/python/__init__.py
index 5470430..b2cc63b 100644
--- a/tensorflow/python/__init__.py
+++ b/tensorflow/python/__init__.py
@@ -78,6 +78,7 @@
# Bring in subpackages.
from tensorflow.python import data
+from tensorflow.python import distribute
from tensorflow.python import keras
from tensorflow.python.feature_column import feature_column_lib as feature_column
from tensorflow.python.layers import layers
@@ -144,21 +145,21 @@
# Export protos
# pylint: disable=undefined-variable
-tf_export('AttrValue')(AttrValue)
-tf_export('ConfigProto')(ConfigProto)
+tf_export(v1=['AttrValue'])(AttrValue)
+tf_export(v1=['ConfigProto'])(ConfigProto)
tf_export('Event', 'summary.Event')(Event)
-tf_export('GPUOptions')(GPUOptions)
-tf_export('GraphDef')(GraphDef)
-tf_export('GraphOptions')(GraphOptions)
-tf_export('HistogramProto')(HistogramProto)
-tf_export('LogMessage')(LogMessage)
-tf_export('MetaGraphDef')(MetaGraphDef)
-tf_export('NameAttrList')(NameAttrList)
-tf_export('NodeDef')(NodeDef)
-tf_export('OptimizerOptions')(OptimizerOptions)
-tf_export('RunMetadata')(RunMetadata)
-tf_export('RunOptions')(RunOptions)
-tf_export('SessionLog', 'summary.SessionLog')(SessionLog)
+tf_export(v1=['GPUOptions'])(GPUOptions)
+tf_export(v1=['GraphDef'])(GraphDef)
+tf_export(v1=['GraphOptions'])(GraphOptions)
+tf_export(v1=['HistogramProto'])(HistogramProto)
+tf_export(v1=['LogMessage'])(LogMessage)
+tf_export(v1=['MetaGraphDef'])(MetaGraphDef)
+tf_export(v1=['NameAttrList'])(NameAttrList)
+tf_export(v1=['NodeDef'])(NodeDef)
+tf_export(v1=['OptimizerOptions'])(OptimizerOptions)
+tf_export(v1=['RunMetadata'])(RunMetadata)
+tf_export(v1=['RunOptions'])(RunOptions)
+tf_export(v1=['SessionLog', 'summary.SessionLog'])(SessionLog)
tf_export('Summary', 'summary.Summary')(Summary)
tf_export('summary.SummaryDescription')(SummaryDescription)
tf_export('SummaryMetadata')(SummaryMetadata)
diff --git a/tensorflow/python/autograph/__init__.py b/tensorflow/python/autograph/__init__.py
index 7252e0d..6faeb01 100644
--- a/tensorflow/python/autograph/__init__.py
+++ b/tensorflow/python/autograph/__init__.py
@@ -12,10 +12,20 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Autograph compiles Python code into equivalent TensorFlow code.
+"""Conversion of plain Python into TensorFlow graph code.
-Equivalent here means that they have the same effect when executed.
+NOTE: In TensorFlow 2.0, AutoGraph is automatically applied when using
+`tf.function`. This module contains lower-level APIs for advanced use.
+
+For more information, see the
+[AutoGraph guide](https://www.tensorflow.org/guide/autograph).
+
+By equivalent graph code we mean code that generates a TensorFlow graph when
+run. The generated graph has the same effects as the original code when executed
+(for example with `tf.function` or `tf.compat.v1.Session.run`). In other words,
+using AutoGraph can be thought of as running Python in TensorFlow.
"""
+# TODO(b/119833526): Link to the new tf.function + autograph tutorial.
from __future__ import absolute_import
from __future__ import division
@@ -43,6 +53,7 @@
from tensorflow.python.autograph.pyct.transformer import AutographParseError
from tensorflow.python.util.all_util import remove_undocumented
+# TODO(mdan): Revisit this list once we finalize the generated code mechanism.
_allowed_symbols = [
# Main API
'ConversionOptions',
diff --git a/tensorflow/python/autograph/converters/call_trees.py b/tensorflow/python/autograph/converters/call_trees.py
index 9b85fc8..d4eb17e 100644
--- a/tensorflow/python/autograph/converters/call_trees.py
+++ b/tensorflow/python/autograph/converters/call_trees.py
@@ -183,7 +183,7 @@
for dec in target_node.decorator_list:
decorator_fn = self._resolve_decorator_name(dec)
if (decorator_fn is not None and
- decorator_fn in self.ctx.program.options.strip_decorators):
+ self.ctx.program.options.should_strip(decorator_fn)):
return False
return True
@@ -261,7 +261,7 @@
func=func,
owner=owner,
options=self.ctx.program.options.to_ast(
- self.ctx.info.namespace,
+ self.ctx,
internal_convert_user_code=self.ctx.program.options.recursive),
args=node.args)
# TODO(mdan): Improve the template mechanism to better support this.
@@ -323,12 +323,12 @@
# 1. super() calls - these are preserved. The class conversion mechanism
# will ensure that they return the correct value.
- if ast_util.matches(node, 'super(_)'):
+ if ast_util.matches(node, parser.parse_expression('super(_)')):
return node
# 2. super().method calls - these are preserved as well, when the
# conversion processes the entire class.
- if (ast_util.matches(node, 'super(_)._(_)') and
+ if (ast_util.matches(node, parser.parse_expression('super(_)._(_)')) and
self.ctx.info.owner_type is not None):
return node
diff --git a/tensorflow/python/autograph/converters/control_flow.py b/tensorflow/python/autograph/converters/control_flow.py
index bef6cae..a39a0b0 100644
--- a/tensorflow/python/autograph/converters/control_flow.py
+++ b/tensorflow/python/autograph/converters/control_flow.py
@@ -49,7 +49,13 @@
def _create_cond_branch(self, body_name, aliased_orig_names,
aliased_new_names, body, returns):
- if len(returns) == 1:
+ if not returns:
+ # TODO(b/110167197): Replace with a plain return.
+ template = """
+ return 1
+ """
+ return_stmt = templates.replace(template)
+ elif len(returns) == 1:
template = """
return retval
"""
@@ -220,7 +226,7 @@
# branch functions will return a dummy value that ensures cond
# actually has some return value as well.
cond_results = None
- # TODO(mdan): This doesn't belong here; it's specific to the operator.
+ # TODO(mdan): Replace with None once side_effect_guards is retired.
returned_from_body = (templates.replace_as_expression(
'ag__.match_staging_level(1, cond_var_name)',
cond_var_name=cond_var_name),)
@@ -278,14 +284,6 @@
' these symbols before the loop'.format(
self._fmt_symbols(live_defs_in_loop)))
- if not loop_state:
- # TODO(mdan): Implement this properly.
- # We need to check whether any variable created inside the body scope
- # is used before being modified outside the scope. This should be done
- # during activity analysis, and in general should cover the case where
- # variables may not be initialized.
- raise ValueError('cannot convert loop: no outputs')
-
return loop_state, reserved_symbols
def _state_constructs(self, loop_state, reserved_symbols):
@@ -337,26 +335,44 @@
node_body = ast_util.rename_symbols(node.body, ssf_map)
test = ast_util.rename_symbols(node.test, ssf_map)
- template = """
- def test_name(state_ssf):
- return test
- def body_name(state_ssf):
- body
- return state_ssf,
- state_ast_tuple = ag__.while_stmt(
- test_name, body_name, (state,), (extra_deps,))
- """
- node = templates.replace(
- template,
- state=loop_state,
- state_ssf=state_ssf,
- state_ast_tuple=state_ast_tuple,
- test_name=self.ctx.namer.new_symbol('loop_test', reserved_symbols),
- test=test,
- body_name=self.ctx.namer.new_symbol('loop_body', reserved_symbols),
- body=node_body,
- extra_deps=tuple(s.ast() for s in cond_closure),
- )
+ if loop_state:
+ template = """
+ def test_name(state_ssf):
+ return test
+ def body_name(state_ssf):
+ body
+ return state_ssf,
+ state_ast_tuple = ag__.while_stmt(
+ test_name, body_name, (state,), (extra_deps,))
+ """
+ node = templates.replace(
+ template,
+ state=loop_state,
+ state_ssf=state_ssf,
+ state_ast_tuple=state_ast_tuple,
+ test_name=self.ctx.namer.new_symbol('loop_test', reserved_symbols),
+ test=test,
+ body_name=self.ctx.namer.new_symbol('loop_body', reserved_symbols),
+ body=node_body,
+ extra_deps=tuple(s.ast() for s in cond_closure),
+ )
+ else:
+ template = """
+ def test_name():
+ return test
+ def body_name():
+ body
+ return ()
+ ag__.while_stmt(test_name, body_name, (), (extra_deps,))
+ """
+ node = templates.replace(
+ template,
+ test_name=self.ctx.namer.new_symbol('loop_test', reserved_symbols),
+ test=test,
+ body_name=self.ctx.namer.new_symbol('loop_body', reserved_symbols),
+ body=node_body,
+ extra_deps=tuple(s.ast() for s in cond_closure),
+ )
return node
@@ -373,29 +389,50 @@
else:
extra_test = parser.parse_expression('True')
- template = """
- def extra_test_name(state_ssf):
- return extra_test_expr
- def body_name(loop_vars, state_ssf):
- # Workaround for PEP-3113
- iterate = loop_vars
- body
- return state_ssf,
- state_ast_tuple = ag__.for_stmt(
- iter_, extra_test_name, body_name, (state,))
- """
- node = templates.replace(
- template,
- state=loop_state,
- state_ssf=state_ssf,
- state_ast_tuple=state_ast_tuple,
- iter_=node.iter,
- iterate=node.target,
- extra_test_name=self.ctx.namer.new_symbol('extra_test',
- reserved_symbols),
- extra_test_expr=extra_test,
- body_name=self.ctx.namer.new_symbol('loop_body', reserved_symbols),
- body=node_body)
+ if loop_state:
+ template = """
+ def extra_test_name(state_ssf):
+ return extra_test_expr
+ def body_name(loop_vars, state_ssf):
+ # Workaround for PEP-3113
+ iterate = loop_vars
+ body
+ return state_ssf,
+ state_ast_tuple = ag__.for_stmt(
+ iter_, extra_test_name, body_name, (state,))
+ """
+ node = templates.replace(
+ template,
+ state=loop_state,
+ state_ssf=state_ssf,
+ state_ast_tuple=state_ast_tuple,
+ iter_=node.iter,
+ iterate=node.target,
+ extra_test_name=self.ctx.namer.new_symbol('extra_test',
+ reserved_symbols),
+ extra_test_expr=extra_test,
+ body_name=self.ctx.namer.new_symbol('loop_body', reserved_symbols),
+ body=node_body)
+ else:
+ template = """
+ def extra_test_name():
+ return extra_test_expr
+ def body_name(loop_vars):
+ # Workaround for PEP-3113
+ iterate = loop_vars
+ body
+ return ()
+ ag__.for_stmt(iter_, extra_test_name, body_name, ())
+ """
+ node = templates.replace(
+ template,
+ iter_=node.iter,
+ iterate=node.target,
+ extra_test_name=self.ctx.namer.new_symbol('extra_test',
+ reserved_symbols),
+ extra_test_expr=extra_test,
+ body_name=self.ctx.namer.new_symbol('loop_body', reserved_symbols),
+ body=node_body)
return node
diff --git a/tensorflow/python/autograph/converters/side_effect_guards.py b/tensorflow/python/autograph/converters/side_effect_guards.py
index 98e29ec..d7c0951 100644
--- a/tensorflow/python/autograph/converters/side_effect_guards.py
+++ b/tensorflow/python/autograph/converters/side_effect_guards.py
@@ -85,11 +85,26 @@
new_alias_map.update(alias_map)
alias_map = new_alias_map
current_dest = new_dest
- if reindent_requested and not current_dest:
- # TODO(mdan): There may still be something that could be done.
- raise ValueError('Unable to insert statement into the computation flow: '
- 'it is not followed by any computation which '
- 'the statement could gate.')
+
+ if reindent_requested:
+ no_controls_to_gate = False
+ if not current_dest:
+ no_controls_to_gate = True
+ if len(current_dest) == 1:
+ if ast_util.matches(current_dest[0], 'return'):
+ no_controls_to_gate = True
+ if ast_util.matches(current_dest[0], 'return ()'):
+ no_controls_to_gate = True
+ if ast_util.matches(current_dest[0], 'return []'):
+ no_controls_to_gate = True
+ if ast_util.matches(current_dest[0], 'return {}'):
+ no_controls_to_gate = True
+ if no_controls_to_gate:
+ # TODO(mdan): There may still be something that could be done.
+ raise ValueError(
+ 'Unable to insert statement into the computation flow: it is not'
+ ' followed by any computation which the statement could gate.')
+
return new_nodes
def visit_FunctionDef(self, node):
diff --git a/tensorflow/python/autograph/core/converter.py b/tensorflow/python/autograph/core/converter.py
index e88c467..4543b11 100644
--- a/tensorflow/python/autograph/core/converter.py
+++ b/tensorflow/python/autograph/core/converter.py
@@ -63,8 +63,9 @@
from __future__ import division
from __future__ import print_function
-from enum import Enum
-from enum import IntEnum
+import weakref
+
+import enum
from tensorflow.python.autograph.core import config
from tensorflow.python.autograph.core import naming
@@ -83,6 +84,7 @@
from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
from tensorflow.python.autograph.pyct.static_analysis import type_info
from tensorflow.python.eager import function
+from tensorflow.python.util.tf_export import tf_export
# TODO(mdan): These contexts can be refactored into first class objects.
# For example, we could define Program and Entity abstractions that hold on
@@ -91,37 +93,42 @@
# TODO(mdan): Add a test specific to this converter.
-class Verbosity(IntEnum):
- """Different levels of verbosity for printing errors.
+@tf_export('autograph.experimental.Verbosity')
+class Verbosity(enum.IntEnum):
+ """Represents conversion verbosity levels.
Attributes:
- * BRIEF: No logging, minimal error messages.
- * VERBOSE: Detailed logging of generated code, detailed error messages.
+ BRIEF: No logging, minimal error messages.
+ VERBOSE: Detailed logging of generated code, detailed error messages.
"""
+
BRIEF = 0
VERBOSE = 1
-class Feature(Enum):
- """Constants to use when selecting AutoGraph features."""
+@tf_export('autograph.experimental.Feature')
+class Feature(enum.Enum):
+ """Represents conversion options that can be toggled on or off.
- ALL = 'Enable all features.'
+ Attributes:
+ ALL: Enable all features.
+ AUTO_CONTROL_DEPS: Insert of control dependencies in the generated code.
+ DECORATORS: Allow decorators in local functions. Note that special
+ decorators, like `tf.function`, are allowed regardless of this toggle.
+ ERROR_REWRITING: Rewrite errors that occur in the generated code to
+ indicate the source code to which the failing code corresponds.
+ LISTS: Convert list idioms, like initializers, slices, append, etc.
+ NAME_SCOPES: Insert name scopes that name ops according to context, like the
+ function they were defined in.
+ """
- AUTO_CONTROL_DEPS = (
- 'Insert of control dependencies in the generated code.')
- DECORATORS = (
- 'Allow decorators in local functions. Note that special decorators,'
- ' like ag.convert or tf.function are allowed regardless of this toggle.')
- ERROR_REWRITING = (
- 'Rewrite errors that occur in the generated code to indicate the source'
- ' code to which the failing code corresponds.')
- LISTS = 'Convert list idioms, like initializers, slices, append, etc.'
- NAME_SCOPES = (
- 'Insert name scopes that name ops according to context, like the'
- ' function they were defined in.')
+ ALL = 'ALL'
- def __repr__(self):
- return self.name
+ AUTO_CONTROL_DEPS = 'AUTO_CONTROL_DEPS'
+ DECORATORS = 'DECORATORS'
+ ERROR_REWRITING = 'ERROR_REWRITING'
+ LISTS = 'LISTS'
+ NAME_SCOPES = 'NAME_SCOPES'
class ConversionOptions(object):
@@ -157,7 +164,9 @@
# TODO(mdan): Rename to conversion_recursion_depth?
self.internal_convert_user_code = internal_convert_user_code
- if isinstance(optional_features, Feature):
+ if optional_features is None:
+ optional_features = ()
+ elif isinstance(optional_features, Feature):
optional_features = (optional_features,)
optional_features = frozenset(optional_features)
self.optional_features = optional_features
@@ -168,19 +177,28 @@
# TODO(mdan): Revert if function.defun becomes a public symbol.
return self._strip_decorators + (function.defun,)
+ def should_strip(self, decorator):
+ for blacklisted in self.strip_decorators:
+ if blacklisted is decorator:
+ return True
+ if isinstance(blacklisted, weakref.ref):
+ blacklisted_deref = blacklisted()
+ if (blacklisted_deref is not None and blacklisted_deref is decorator):
+ return True
+ return False
+
def uses(self, feature):
return (Feature.ALL in self.optional_features or
feature in self.optional_features)
- def to_ast(self, namespace, internal_convert_user_code=None):
+ def to_ast(self, ctx, internal_convert_user_code=None):
"""Returns a representation of this object as an AST node.
The AST node encodes a constructor that would create an object with the
same contents.
Args:
- namespace: Dict[str, Any], the namespace to use when serializing values to
- names.
+ ctx: EntityContext, the entity with which this AST needs to be consistent.
internal_convert_user_code: Optional[bool], allows ovrriding the
corresponding value.
@@ -198,10 +216,11 @@
"""
def as_qualified_name(o):
- name = inspect_utils.getqualifiedname(namespace, o)
+ name = inspect_utils.getqualifiedname(ctx.info.namespace, o, max_depth=1)
if not name:
- raise ValueError('Could not locate entity {} in {}'.format(
- o, namespace))
+ # TODO(mdan): This needs to account for the symbols defined locally.
+ name = ctx.namer.new_symbol(o.__name__, ())
+ ctx.program.add_symbol(name, weakref.ref(o))
return name
def list_of_names(values):
@@ -272,6 +291,7 @@
self.dependency_cache = {}
self.additional_imports = set()
self.name_map = {}
+ self.additional_symbols = {}
@property
def required_imports(self):
@@ -314,6 +334,11 @@
else:
self.name_map[o] = name
+ def add_symbol(self, name, value):
+ if name in self.additional_symbols:
+ assert self.additional_symbols[name] is value
+ self.additional_symbols[name] = value
+
def add_to_cache(self, original_entity, converted_ast):
self.conversion_order.append(original_entity)
self.dependency_cache[original_entity] = converted_ast
@@ -419,7 +444,7 @@
self.directives = {}
-class AgAnno(Enum):
+class AgAnno(enum.Enum):
"""Annotation labels specific to AutoGraph. See anno.py."""
DIRECTIVES = 'User directives associated with the annotated statement.'
diff --git a/tensorflow/python/autograph/core/converter_test.py b/tensorflow/python/autograph/core/converter_test.py
index b73c67e..864ea6c 100644
--- a/tensorflow/python/autograph/core/converter_test.py
+++ b/tensorflow/python/autograph/core/converter_test.py
@@ -18,6 +18,8 @@
from __future__ import division
from __future__ import print_function
+import weakref
+
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.core import converter_testing
from tensorflow.python.autograph.pyct import anno
@@ -29,6 +31,36 @@
pass
+class ConversionOptionsTest(test.TestCase):
+
+ def test_should_strip_weakrefs(self):
+ def test_fn():
+ pass
+
+ def weak_test_fn_a():
+ pass
+
+ def weak_test_fn_b():
+ pass
+
+ def weak_test_fn_c():
+ pass
+
+ wr_a = weakref.ref(weak_test_fn_a)
+ # Create an extra weakref to check whether the existence of multiple weak
+ # references influences the process.
+ _ = weakref.ref(weak_test_fn_b)
+ wr_b = weakref.ref(weak_test_fn_b)
+ _ = weakref.ref(weak_test_fn_c)
+
+ opts = converter.ConversionOptions(strip_decorators=(test_fn, wr_a, wr_b))
+
+ self.assertTrue(opts.should_strip(test_fn))
+ self.assertTrue(opts.should_strip(weak_test_fn_a))
+ self.assertTrue(opts.should_strip(weak_test_fn_b))
+ self.assertFalse(opts.should_strip(weak_test_fn_c))
+
+
class ConverterBaseTest(converter_testing.TestCase):
def test_get_definition_directive_basic(self):
diff --git a/tensorflow/python/autograph/impl/api.py b/tensorflow/python/autograph/impl/api.py
index 19a4720..c113f0e 100644
--- a/tensorflow/python/autograph/impl/api.py
+++ b/tensorflow/python/autograph/impl/api.py
@@ -40,6 +40,7 @@
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import tf_decorator
from tensorflow.python.util import tf_inspect
+from tensorflow.python.util.tf_export import tf_export
# TODO(mdan): Properly document the type hints.
# TODO(mdan): Reduce the type hint information to (module, type).
@@ -49,7 +50,10 @@
# TODO(mdan): This should behave like to_graph (e.g. convert statically).
# TODO(znado): Make an alias so can write Verbosity directly without needing
# to write converter.
-def convert(recursive=False, verbose=converter.Verbosity.VERBOSE):
+def convert(
+ recursive=False,
+ verbose=converter.Verbosity.BRIEF,
+ optional_features=converter.Feature.ALL):
"""Decorator that compiles a function to use TensorFlow ops.
The decorator is dynamic - it recompiles the target whenever the decorated
@@ -61,6 +65,9 @@
recursive: bool, whether to recursively convert any functions or classes
that the converted function may use.
verbose: converter.Verbosity, the level of verbosity.
+ optional_features: converted.Feature, allows toggling optional or
+ experimental features. When set to None, only the core features are
+ enabled.
Returns:
Callable, a decorator that converts the given function into an equivalent
@@ -78,7 +85,7 @@
recursive=recursive,
verbose=verbose,
force_conversion=True,
- optional_features=converter.Feature.ALL,
+ optional_features=optional_features,
), *args, **kwargs)
wrapper = tf_decorator.make_decorator(f, wrapper)
@@ -151,11 +158,9 @@
return decorator
-# TODO(mdan): Move to a private, undocumented module.
def converted_call(f, owner, options, *args, **kwargs):
"""Compiles a function call inline. For internal use only."""
- if options.verbose >= converter.Verbosity.VERBOSE:
- logging.info('Converted call: {}; owner: {}'.format(f, owner))
+ logging.vlog(logging.DEBUG, 'Converted call: %s; owner: %s', f, owner)
if owner is not None:
if not isinstance(f, str):
@@ -196,7 +201,7 @@
return f(*args, **kwargs)
# Unwrap functools.partial objects
- # TODO(allenl, mdan): Consider sharing unwrapping logic with tf_inspect.
+ # TODO(mdan): Consider sharing unwrapping logic with tf_inspect.
while isinstance(f, functools.partial):
args = f.args + args
new_kwargs = {}
@@ -277,9 +282,9 @@
verbose=options.verbose,
arg_values=arg_values,
arg_types=arg_types,
- partial_types=partial_types,
strip_decorators=options.strip_decorators,
- optional_features=options.optional_features)
+ optional_features=options.optional_features,
+ experimental_partial_types=partial_types)
result = converted_f(*effective_args, **kwargs)
@@ -308,44 +313,81 @@
return False
-# TODO(mdan): Rename: to_ops?
-# TODO(mdan): Look into overloading as function and decorator, like tfe.defun?
-# TODO(mdan): Remove partial_types.
-def to_graph(e,
+@tf_export('autograph.to_graph')
+def to_graph(entity,
recursive=True,
verbose=converter.Verbosity.VERBOSE,
arg_values=None,
arg_types=None,
- partial_types=None,
strip_decorators=None,
- optional_features=converter.Feature.ALL):
- """Converts a Python entity into equivalent code that uses TensorFlow ops.
+ optional_features=converter.Feature.ALL,
+ experimental_partial_types=None):
+ """Converts a Python entity into a TensorFlow graph.
+
+ Also see: `tf.autograph.to_code`, `tf.function`.
+
+ Unlike `tf.function`, `to_graph` is a low-level transpiler that converts
+ Python code to TensorFlow graph code. It does not implement any caching,
+ variable management or create any actual ops, and is best used where greater
+ control over the generated TensorFlow graph is desired. Another difference
+ from `tf.function` is that `to_graph` will not wrap the graph into a
+ TensorFlow function or a Python callable. Internally, `tf.function` uses
+ `to_graph`.
+
+ _Example Usage_
+
+ ```python
+ def foo(x):
+ if x > 0:
+ y = x * x
+ else:
+ y = -x
+ return y
+
+ converted_foo = to_graph(foo)
+
+ x = tf.constant(1)
+ y = converted_foo(x) # converted_foo is a TensorFlow Op-like.
+ assert is_tensor(y)
+ ```
Supported Python entities include:
* functions
* classes
+ * object methods
- Classes are converted by converting all their methods into a new class.
+ Functions are converted into new functions with converted code.
+
+ Classes are converted by generating a new class whose methods use converted
+ code.
+
+ Methods are converted into unbound function that have an additional first
+ argument called `self`.
Args:
- e: Union[Callable, Type], the Python entity to convert.
- recursive: bool, whether to recursively convert any functions that the
+ entity: Python callable or class to convert.
+ recursive: Whether to recursively convert any functions that the
converted function may call.
- verbose: converter.Verbosity, the level of printing verbosity to use.
- arg_values: Optional[Dict[Text, Any]], value hints for symbols including
- function arguments.
- arg_types: Optional[Dict[Text, Type]], type hints for symbols including
- function arguments.
- partial_types: Set[Type], reserved for internal use.
- strip_decorators: Tuple[Callable], same as
- ConversionOptions.strip_decorators.
- optional_features: Union[Feature, Set[Feature]], same as
- ConversionOptions.optional_features.
+ verbose: The level of printing verbosity to use, as a
+ `tf.autograph.experimental.Verbosity` value.
+ arg_values: Optional dict of value hints for symbols including
+ function arguments mapping string names to actual values. For example,
+ `arg_values={'a': 1}` will map the variable `a` to the value `1`.
+ arg_types: Optional dict of type hints for symbols including function
+ arguments. Type hints allow specifying just the type of a variable, rather
+ than a specific value.
+ strip_decorators: A tuple specifying decorators that should be
+ excluded from the compiled output. By default, when converting a function
+ before the decorators are applied, the compiled output will include those
+ decorators.
+ optional_features: `None`, a tuple of, or a single
+ `tf.autograph.experimental.Feature` value. Controls the use of
+ optional features in the conversion process.
+ experimental_partial_types: A `set` of `type` values, reserved for internal
+ use.
Returns:
- Union[Callable, Type], the converted entity, which is the same kind as e
- (that is, a function is e is a function, a class if e is a class, etc.) but
- its code has been converted to use TF ops.
+ Same as `entity`, the converted Python function or class.
Raises:
ValueError: If the entity could not be converted.
@@ -360,11 +402,11 @@
verbose=verbose,
strip_decorators=strip_decorators,
optional_features=optional_features),
- partial_types=partial_types,
+ partial_types=experimental_partial_types,
autograph_module=tf_inspect.getmodule(to_graph),
uncompiled_modules=config.DEFAULT_UNCOMPILED_MODULES)
- _, name, namespace = conversion.entity_to_graph(e, program_ctx, arg_values,
- arg_types)
+ _, name, namespace = conversion.entity_to_graph(entity, program_ctx,
+ arg_values, arg_types)
nodes = []
for dep in reversed(program_ctx.conversion_order):
@@ -381,10 +423,13 @@
# Avoid overwriting entities that have been transformed.
if key not in compiled_module.__dict__:
compiled_module.__dict__[key] = val
+ for key, val in program_ctx.additional_symbols.items():
+ if key not in compiled_module.__dict__:
+ compiled_module.__dict__[key] = val
compiled = getattr(compiled_module, name)
- if tf_inspect.isfunction(e):
- compiled.__defaults__ = e.__defaults__
+ if tf_inspect.isfunction(entity):
+ compiled.__defaults__ = entity.__defaults__
if hasattr(compiled, '__globals__'):
# Remove self to avoid circular references. This will probably only work
@@ -409,38 +454,52 @@
return compiled
-def to_code(e,
+@tf_export('autograph.to_code')
+def to_code(entity,
recursive=True,
arg_values=None,
arg_types=None,
- partial_types=None,
- indentation=' '):
- """Returns the equivalent code that uses TensorFlow ops.
+ indentation=' ',
+ optional_features=converter.Feature.ALL,
+ experimental_partial_types=None):
+ """Similar to `to_graph`, but returns Python source code as a string.
- Also see: `to_graph`, `convert`
+ Also see: `tf.autograph.to_graph`.
+
+ `to_graph` returns the Python source code that can be used to generate a
+ TensorFlow graph that is functionally identical to the input Python code.
Args:
- e: Union[Callable, Type], the Python entity to convert.
- recursive: bool, whether to recursively convert any functions that the
+ entity: Python callable or class to convert.
+ recursive: Whether to recursively convert any functions that the
converted function may call.
- arg_values: Optional[Dict[Text, Any]], value hints for symbols including
- function arguments.
- arg_types: Optional[Dict[Text, Type]], type hints for symbols including
- function arguments.
- partial_types: Set[Type], reserved for internal use.
- indentation: Text, when to use for each level of indentation.
+ arg_values: Optional dict of value hints for symbols including
+ function arguments mapping string names to actual values. For example,
+ `arg_values={'a': 1}` will map the variable `a` to the value `1`.
+ arg_types: Optional dict of type hints for symbols including function
+ arguments. Type hints allow specifying just the type of a variable, rather
+ than a specific value.
+ indentation: The string to use for indenting. Typically two or four spaces,
+ or just the tab character.
+ optional_features: `None`, a tuple of, or a single
+ `tf.autograph.experimental.Feature` value. Controls the use of
+ optional features in the conversion process.
+ experimental_partial_types: A `set` of `type` values, reserved for internal
+ use.
Returns:
- Text, the converted code.
+ The converted code as string.
"""
program_ctx = converter.ProgramContext(
options=converter.ConversionOptions(
recursive=recursive,
- strip_decorators=(convert, do_not_convert, converted_call)),
- partial_types=partial_types,
+ verbose=converter.Verbosity.BRIEF,
+ strip_decorators=(convert, do_not_convert, converted_call),
+ optional_features=optional_features),
+ partial_types=experimental_partial_types,
autograph_module=tf_inspect.getmodule(to_graph),
uncompiled_modules=config.DEFAULT_UNCOMPILED_MODULES)
- conversion.entity_to_graph(e, program_ctx, arg_values, arg_types)
+ conversion.entity_to_graph(entity, program_ctx, arg_values, arg_types)
code = '\n'.join(
compiler.ast_to_source(program_ctx.dependency_cache[dep], indentation)
diff --git a/tensorflow/python/autograph/impl/api_test.py b/tensorflow/python/autograph/impl/api_test.py
index 66edda5..d5561ba 100644
--- a/tensorflow/python/autograph/impl/api_test.py
+++ b/tensorflow/python/autograph/impl/api_test.py
@@ -218,6 +218,7 @@
constant_op.constant(-1))
self.assertEqual(1, self.evaluate(x))
+ @test_util.run_v1_only('b/120545219')
def test_converted_call_functools_partial(self):
def test_fn(x, y, z):
diff --git a/tensorflow/python/autograph/impl/conversion.py b/tensorflow/python/autograph/impl/conversion.py
index f8decd2..d0a12df 100644
--- a/tensorflow/python/autograph/impl/conversion.py
+++ b/tensorflow/python/autograph/impl/conversion.py
@@ -129,8 +129,7 @@
Raises:
ValueError: if the entity type is not supported.
"""
- if program_ctx.options.verbose == converter.Verbosity.VERBOSE:
- logging.info('Converting {}'.format(o))
+ logging.vlog(logging.DEBUG, 'Converting %s', o)
if tf_inspect.isclass(o):
node, name, ns = class_to_graph(o, program_ctx)
@@ -164,9 +163,9 @@
program_ctx.add_to_cache(o, node)
- if program_ctx.options.verbose == converter.Verbosity.VERBOSE:
- logging.info('Compiled output of {}:\n\n{}\n'.format(
- o, compiler.ast_to_source(node)))
+ if logging.get_verbosity() <= logging.DEBUG:
+ logging.vlog(logging.DEBUG, 'Compiled output of %s:\n\n%s\n', o,
+ compiler.ast_to_source(node))
if program_ctx.options.recursive:
while True:
diff --git a/tensorflow/python/autograph/operators/control_flow.py b/tensorflow/python/autograph/operators/control_flow.py
index 6708977..afa3787 100644
--- a/tensorflow/python/autograph/operators/control_flow.py
+++ b/tensorflow/python/autograph/operators/control_flow.py
@@ -20,7 +20,6 @@
from tensorflow.python.autograph.operators import py_builtins
from tensorflow.python.data.ops import dataset_ops
-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 gen_math_ops
@@ -88,7 +87,10 @@
def while_body(iterate_index, *state):
iterate = iter_[iterate_index]
new_state = body(iterate, *state)
- return (iterate_index + 1,) + new_state
+ if new_state:
+ return (iterate_index + 1,) + new_state
+ else:
+ return iterate_index + 1
def while_cond(iterate_index, *state):
return gen_math_ops.logical_and(iterate_index < n, extra_test(*state))
@@ -99,51 +101,33 @@
init_state=(0,) + init_state,
extra_deps=(iter_,),
opts=dict(maximum_iterations=n))
- # Dropping the iteration index because it's not syntactically visible.
- results = results[1:]
- # TODO(mdan): Remove this special case.
- if len(results) == 1:
- return results[0]
+ # Dropping the iteration index because it's not syntactically visible.
+ # TODO(mdan): Don't.
+ if isinstance(results, (tuple, list)):
+ assert len(results) >= 1 # Has at least the iterate.
+ if len(results) > 1:
+ results = results[1:]
+ if len(results) == 1:
+ # TODO(mdan): Remove this special case.
+ results, = results
+ else:
+ results = ()
+
return results
def _dataset_for_stmt(ds, extra_test, body, init_state):
"""Overload of for_stmt that iterates over TF Datasets."""
- # Because Datsets only expose get_next, in the style of Python iterators,
- # we are forced to unpack the loop as:
- #
- # epoch_number, iterate = ds.get_next()
- # while epoch_number < 2:
- # <body>
- # epoch_number, iterate = ds.get_next()
- epoch_numbers = dataset_ops.Dataset.range(2)
- def tag_with(ds, tag):
- return dataset_ops.Dataset.zip(
- (dataset_ops.Dataset.from_tensors(tag).repeat(), ds))
- ds_with_epoch = epoch_numbers.flat_map(lambda i: tag_with(ds, i))
+ if extra_test(*init_state) is not True:
+ raise NotImplementedError(
+ 'break statements are not yet supported in for/Dataset loops')
- iterator = dataset_ops.make_initializable_iterator(ds_with_epoch)
- with ops.control_dependencies((iterator.initializer,)):
- epoch_number, iterate = iterator.get_next()
+ def reduce_body(state, iterate):
+ new_state = body(iterate, *state)
+ return new_state
- def while_body(epoch_number, iterate, *state):
- new_state = body(iterate, *state)
- epoch_number, iterate = iterator.get_next()
- return (epoch_number, iterate) + new_state
-
- def while_cond(epoch_number, iterate, *state):
- del iterate
- return gen_math_ops.logical_and(epoch_number < 1, extra_test(*state))
-
- results = while_stmt(
- while_cond,
- while_body,
- init_state=(epoch_number, iterate) + init_state,
- extra_deps=())
- # Dropping the epoch number and iterate because they are not syntactically
- # visible.
- results = results[2:]
+ results = ds.reduce(init_state, reduce_body)
# TODO(mdan): Remove this special case.
if len(results) == 1:
diff --git a/tensorflow/python/autograph/operators/py_builtins.py b/tensorflow/python/autograph/operators/py_builtins.py
index 2f55d53..ddf05f7 100644
--- a/tensorflow/python/autograph/operators/py_builtins.py
+++ b/tensorflow/python/autograph/operators/py_builtins.py
@@ -174,6 +174,7 @@
override_kwargs['flush'] = True
def print_wrapper(*vals):
+ vals = tuple(v.numpy() if tensor_util.is_tensor(v) else v for v in vals)
if six.PY3:
# TensorFlow doesn't seem to generate Unicode when passing strings to
# py_func. This causes the print to add a "b'" wrapper to the output,
@@ -193,6 +194,7 @@
def _tf_range(start_or_stop, stop, step):
+ """Overload of range_ that generates a TF range tensor."""
# Note: for static inputs (e.g. constants), tf.range errors out at graph
# construction time, instead of returning an empty tensor. Preventing the
# graph construction error aligns the semantics with Python.
diff --git a/tensorflow/python/autograph/pyct/ast_util.py b/tensorflow/python/autograph/pyct/ast_util.py
index ea7eca6..3dc10cf 100644
--- a/tensorflow/python/autograph/pyct/ast_util.py
+++ b/tensorflow/python/autograph/pyct/ast_util.py
@@ -200,7 +200,8 @@
bool
"""
if isinstance(pattern, str):
- pattern = parser.parse_expression(pattern)
+ pattern, = parser.parse_str(pattern).body
+
matcher = PatternMatcher(pattern)
matcher.visit(node)
return matcher.matches
diff --git a/tensorflow/python/autograph/pyct/compiler.py b/tensorflow/python/autograph/pyct/compiler.py
index 06e66c5..aa4fd55 100644
--- a/tensorflow/python/autograph/pyct/compiler.py
+++ b/tensorflow/python/autograph/pyct/compiler.py
@@ -67,6 +67,12 @@
trimmed_code_lines.append(l)
code = '\n'.join(trimmed_code_lines)
+ # Work around the reference cycle generated by astor.
+ # See https://github.com/berkerpeksag/astor/blob/55dd323f7d8d696610c703c0296763c567685c31/astor/code_gen.py#L162 # pylint:disable=line-too-long
+ # Reference cycles are quite disliked by TensorFlow's tests.
+ if hasattr(generator, 'write'):
+ generator.write = None
+
return code
diff --git a/tensorflow/python/autograph/pyct/inspect_utils.py b/tensorflow/python/autograph/pyct/inspect_utils.py
index 7c819f3..360dd83 100644
--- a/tensorflow/python/autograph/pyct/inspect_utils.py
+++ b/tensorflow/python/autograph/pyct/inspect_utils.py
@@ -101,7 +101,7 @@
return namespace
-def getqualifiedname(namespace, object_, max_depth=2):
+def getqualifiedname(namespace, object_, max_depth=5, visited=None):
"""Returns the name by which a value can be referred to in a given namespace.
If the object defines a parent module, the function attempts to use it to
@@ -115,16 +115,20 @@
object_: Any, the value to search.
max_depth: Optional[int], a limit to the recursion depth when searching
inside modules.
+ visited: Optional[Set[int]], ID of modules to avoid visiting.
Returns: Union[str, None], the fully-qualified name that resolves to the value
o, or None if it couldn't be found.
"""
- for name, value in namespace.items():
+ if visited is None:
+ visited = set()
+
+ for name in namespace:
# The value may be referenced by more than one symbol, case in which
# any symbol will be fine. If the program contains symbol aliases that
# change over time, this may capture a symbol that will later point to
# something else.
# TODO(mdan): Prefer the symbol that matches the value type name.
- if object_ is value:
+ if object_ is namespace[name]:
return name
# If an object is not found, try to search its parent modules.
@@ -132,22 +136,25 @@
if (parent is not None and parent is not object_ and
parent is not namespace):
# No limit to recursion depth because of the guard above.
- parent_name = getqualifiedname(namespace, parent, max_depth=0)
+ parent_name = getqualifiedname(
+ namespace, parent, max_depth=0, visited=visited)
if parent_name is not None:
- name_in_parent = getqualifiedname(parent.__dict__, object_, max_depth=0)
+ name_in_parent = getqualifiedname(
+ parent.__dict__, object_, max_depth=0, visited=visited)
assert name_in_parent is not None, (
'An object should always be found in its owner module')
return '{}.{}'.format(parent_name, name_in_parent)
- # TODO(mdan): Use breadth-first search and avoid visiting modules twice.
if max_depth:
# Iterating over a copy prevents "changed size due to iteration" errors.
# It's unclear why those occur - suspecting new modules may load during
# iteration.
- for name, value in namespace.copy().items():
- if tf_inspect.ismodule(value):
+ for name in namespace.keys():
+ value = namespace[name]
+ if tf_inspect.ismodule(value) and id(value) not in visited:
+ visited.add(id(value))
name_in_module = getqualifiedname(value.__dict__, object_,
- max_depth - 1)
+ max_depth - 1, visited)
if name_in_module is not None:
return '{}.{}'.format(name, name_in_module)
return None
diff --git a/tensorflow/python/autograph/pyct/inspect_utils_test.py b/tensorflow/python/autograph/pyct/inspect_utils_test.py
index a2c3905..420a20c 100644
--- a/tensorflow/python/autograph/pyct/inspect_utils_test.py
+++ b/tensorflow/python/autograph/pyct/inspect_utils_test.py
@@ -183,6 +183,63 @@
self.assertEqual(inspect_utils.getqualifiedname(ns, bar), 'bar')
self.assertEqual(inspect_utils.getqualifiedname(ns, baz), 'bar.baz')
+ def test_getqualifiedname_efficiency(self):
+ foo = object()
+
+ # We create a densely connected graph consisting of a relatively small
+ # number of modules and hide our symbol in one of them. The path to the
+ # symbol is at least 10, and each node has about 10 neighbors. However,
+ # by skipping visited modules, the search should take much less.
+ ns = {}
+ prev_level = []
+ for i in range(10):
+ current_level = []
+ for j in range(10):
+ mod_name = 'mod_{}_{}'.format(i, j)
+ mod = imp.new_module(mod_name)
+ current_level.append(mod)
+ if i == 9 and j == 9:
+ mod.foo = foo
+ if prev_level:
+ # All modules at level i refer to all modules at level i+1
+ for prev in prev_level:
+ for mod in current_level:
+ prev.__dict__[mod.__name__] = mod
+ else:
+ for mod in current_level:
+ ns[mod.__name__] = mod
+ prev_level = current_level
+
+ self.assertIsNone(inspect_utils.getqualifiedname(ns, inspect_utils))
+ self.assertIsNotNone(
+ inspect_utils.getqualifiedname(ns, foo, max_depth=10000000000))
+
+ def test_getqualifiedname_cycles(self):
+ foo = object()
+
+ # We create a graph of modules that contains circular references. The
+ # search process should avoid them. The searched object is hidden at the
+ # bottom of a path of length roughly 10.
+ ns = {}
+ mods = []
+ for i in range(10):
+ mod = imp.new_module('mod_{}'.format(i))
+ if i == 9:
+ mod.foo = foo
+ # Module i refers to module i+1
+ if mods:
+ mods[-1].__dict__[mod.__name__] = mod
+ else:
+ ns[mod.__name__] = mod
+ # Module i refers to all modules j < i.
+ for prev in mods:
+ mod.__dict__[prev.__name__] = prev
+ mods.append(mod)
+
+ self.assertIsNone(inspect_utils.getqualifiedname(ns, inspect_utils))
+ self.assertIsNotNone(
+ inspect_utils.getqualifiedname(ns, foo, max_depth=10000000000))
+
def test_getqualifiedname_finds_via_parent_module(self):
# TODO(mdan): This test is vulnerable to change in the lib module.
# A better way to forge modules should be found.
diff --git a/tensorflow/python/autograph/pyct/parser.py b/tensorflow/python/autograph/pyct/parser.py
index 8f4037c..d04a401 100644
--- a/tensorflow/python/autograph/pyct/parser.py
+++ b/tensorflow/python/autograph/pyct/parser.py
@@ -24,6 +24,7 @@
import textwrap
import gast
+import six
from tensorflow.python.util import tf_inspect
@@ -91,7 +92,17 @@
def parse_str(src):
"""Returns the AST of given piece of code."""
# TODO(mdan): This should exclude the module things are autowrapped in.
- return gast.parse(src)
+
+ if six.PY2 and '.print(' in src:
+ # This special treatment is required because gast.parse is not aware of
+ # whether print_function was present in the original context.
+ src = 'from __future__ import print_function\n' + src
+ parsed_module = gast.parse(src)
+ parsed_module.body = parsed_module.body[1:]
+ else:
+ parsed_module = gast.parse(src)
+
+ return parsed_module
def parse_expression(src):
@@ -106,7 +117,7 @@
"""
node = parse_str(src)
assert isinstance(node, gast.Module)
- if len(node.body) != 1 and not isinstance(node.body[0], gast.Expr):
+ if len(node.body) != 1 or not isinstance(node.body[0], gast.Expr):
raise ValueError(
'Expected a single expression, found instead %s' % node.body)
return node.body[0].value
diff --git a/tensorflow/python/autograph/pyct/templates.py b/tensorflow/python/autograph/pyct/templates.py
index 9901295..43279b3 100644
--- a/tensorflow/python/autograph/pyct/templates.py
+++ b/tensorflow/python/autograph/pyct/templates.py
@@ -32,6 +32,66 @@
from tensorflow.python.autograph.pyct import qual_names
+class ContextAdjuster(gast.NodeTransformer):
+ """Adjusts the ctx field of nodes to ensure consistency.
+
+ This transformer can change the ctx fields of a variable, tuple and other
+ AST elements that allow one, based on whether the element is being read or
+ written.
+ """
+
+ def __init__(self, override_value):
+ self._ctx_override = override_value
+
+ def visit(self, node):
+ original_override = self._ctx_override
+ node = super(ContextAdjuster, self).visit(node)
+ if hasattr(node, 'ctx'):
+ assert node.ctx is not None, 'node {} has ctx unset'.format(node)
+ self._ctx_override = original_override
+ return node
+
+ def _apply_override(self, node):
+ if self._ctx_override is not None:
+ node.ctx = self._ctx_override()
+
+ def visit_Attribute(self, node):
+ self._apply_override(node)
+ self._ctx_override = gast.Load
+ node = self.generic_visit(node)
+ return node
+
+ def visit_Tuple(self, node):
+ self._apply_override(node)
+ return self.generic_visit(node)
+
+ def visit_List(self, node):
+ self._apply_override(node)
+ return self.generic_visit(node)
+
+ def visit_Name(self, node):
+ self._apply_override(node)
+ return self.generic_visit(node)
+
+ def visit_Call(self, node):
+ self._apply_override(node)
+ # We may be able to override these to Load(), but for now it's simpler
+ # to just assert that they're set.
+ self._ctx_override = None
+ return self.generic_visit(node)
+
+ def visit_Dict(self, node):
+ # We may be able to override these to Load(), but for now it's simpler
+ # to just assert that they're set.
+ self._ctx_override = None
+ return self.generic_visit(node)
+
+ def visit_Subscript(self, node):
+ node.value = self.visit(node.value)
+ self._ctx_override = None
+ return self.generic_visit(node)
+
+
class ReplaceTransformer(gast.NodeTransformer):
"""Replace AST nodes."""
@@ -106,91 +166,6 @@
node.name = repl.id
return node
- def _check_has_context(self, node):
- if not node.ctx:
- raise ValueError('node %s is missing ctx value' % node)
-
- # TODO(mdan): Rewrite _check and _set using a separate transformer.
- def _check_inner_children_have_context(self, node):
- if isinstance(node, gast.Attribute):
- self._check_inner_children_have_context(node.value)
- self._check_has_context(node)
- elif isinstance(node, (gast.Tuple, gast.List)):
- for e in node.elts:
- self._check_inner_children_have_context(e)
- self._check_has_context(node)
- elif isinstance(node, gast.Dict):
- for e in node.keys:
- self._check_inner_children_have_context(e)
- for e in node.values:
- self._check_inner_children_have_context(e)
- elif isinstance(node, gast.Index):
- self._check_inner_children_have_context(node.value)
- elif isinstance(node, gast.Subscript):
- self._check_inner_children_have_context(node.value)
- self._check_inner_children_have_context(node.slice)
- elif isinstance(node, gast.Slice):
- self._check_inner_children_have_context(node.lower)
- if node.upper:
- self._check_inner_children_have_context(node.upper)
- if node.step:
- self._check_inner_children_have_context(node.step)
- elif isinstance(node, gast.BinOp):
- self._check_inner_children_have_context(node.left)
- self._check_inner_children_have_context(node.right)
- elif isinstance(node, gast.UnaryOp):
- self._check_inner_children_have_context(node.operand)
- elif isinstance(node, gast.Name):
- self._check_has_context(node)
- elif isinstance(node, (gast.Str, gast.Num)):
- pass
- elif isinstance(node, gast.Call):
- self._check_inner_children_have_context(node.func)
- for a in node.args:
- self._check_inner_children_have_context(a)
- for k in node.keywords:
- self._check_inner_children_have_context(k.value)
- else:
- raise ValueError('unexpected node type "%s"' % node)
-
- def _set_inner_child_context(self, node, ctx):
- if isinstance(node, gast.Attribute):
- self._set_inner_child_context(node.value, gast.Load())
- node.ctx = ctx
- elif isinstance(node, (gast.Tuple, gast.List)):
- for e in node.elts:
- self._set_inner_child_context(e, ctx)
- node.ctx = ctx
- elif isinstance(node, gast.Name):
- node.ctx = ctx
- elif isinstance(node, gast.Call):
- self._set_inner_child_context(node.func, ctx)
- # We may be able to override these to Load(), but for now it's simpler
- # to just assert that they're set.
- for a in node.args:
- self._check_inner_children_have_context(a)
- for k in node.keywords:
- self._check_inner_children_have_context(k.value)
- elif isinstance(node, gast.Dict):
- # We may be able to override these to Load(), but for now it's simpler
- # to just assert that they're set.
- for e in node.keys:
- self._check_inner_children_have_context(e)
- for e in node.values:
- self._check_inner_children_have_context(e)
- elif isinstance(node, gast.Subscript):
- self._set_inner_child_context(node.value, ctx)
- self._check_inner_children_have_context(node.slice)
- elif isinstance(node, gast.BinOp):
- self._check_inner_children_have_context(node.left)
- self._check_inner_children_have_context(node.right)
- elif isinstance(node, gast.UnaryOp):
- self._check_inner_children_have_context(node.operand)
- elif isinstance(node, (gast.Str, gast.Num)):
- pass
- else:
- raise ValueError('unexpected node type "%s"' % node)
-
def visit_Attribute(self, node):
node = self.generic_visit(node)
if node.attr not in self.replacements:
@@ -209,17 +184,14 @@
new_nodes = self._prepare_replacement(node, node.id)
+ if not new_nodes:
+ return new_nodes
+
# Preserve the target context.
+ adjuster = ContextAdjuster(type(node.ctx))
for n in new_nodes:
- if isinstance(n, (gast.Tuple, gast.List)):
- for e in n.elts:
- self._set_inner_child_context(e, node.ctx)
- if isinstance(n, gast.Attribute):
- # For attributes, the inner Name node receives the context, while the
- # outer ones have it set to Load.
- self._set_inner_child_context(n, node.ctx)
- else:
- n.ctx = node.ctx
+ if hasattr(n, 'ctx'):
+ adjuster.visit(n)
if len(new_nodes) == 1:
new_nodes, = new_nodes
diff --git a/tensorflow/python/autograph/pyct/templates_test.py b/tensorflow/python/autograph/pyct/templates_test.py
index 54019ef..cdb44b8 100644
--- a/tensorflow/python/autograph/pyct/templates_test.py
+++ b/tensorflow/python/autograph/pyct/templates_test.py
@@ -134,19 +134,18 @@
def test_replace_expression_context(self):
template = """
- def test_fn(foo):
+ def test_fn():
foo
"""
node = templates.replace(
template, foo=parser.parse_expression('a + 2 * b / -c'))[0]
- self.assertIsInstance(node.body[0].ctx, gast.Load)
self.assertIsInstance(node.body[0].left.ctx, gast.Load)
self.assertIsInstance(node.body[0].right.left.right.ctx, gast.Load)
def test_replace_complex_context(self):
template = """
- def test_fn(foo):
+ def test_fn():
foo = 0
"""
@@ -160,7 +159,7 @@
def test_replace_index(self):
template = """
- def test_fn(foo):
+ def test_fn():
foo = 0
"""
diff --git a/tensorflow/python/autograph/utils/py_func.py b/tensorflow/python/autograph/utils/py_func.py
index 11ebfb2..ee8b46b 100644
--- a/tensorflow/python/autograph/utils/py_func.py
+++ b/tensorflow/python/autograph/utils/py_func.py
@@ -127,5 +127,6 @@
retval = f(*f_args, **f_kwargs)
return 1 if use_dummy_return else retval
- return script_ops.py_func(f_wrapper, tensor_args, dtypes.int64
- if use_dummy_return else return_dtypes)
+ if use_dummy_return:
+ return_dtypes = dtypes.int32
+ return script_ops.eager_py_func(f_wrapper, tensor_args, return_dtypes)
diff --git a/tensorflow/python/autograph/utils/py_func_test.py b/tensorflow/python/autograph/utils/py_func_test.py
index 28cefd8..d17ede7 100644
--- a/tensorflow/python/autograph/utils/py_func_test.py
+++ b/tensorflow/python/autograph/utils/py_func_test.py
@@ -32,13 +32,13 @@
return a + b + c
with self.cached_session() as sess:
- result = py_func.wrap_py_func(test_fn, dtypes.int64,
+ result = py_func.wrap_py_func(test_fn, dtypes.int32,
(1, constant_op.constant(1), 1))
self.assertEqual(3, self.evaluate(result))
- result = py_func.wrap_py_func(test_fn, dtypes.int64, (1, 1, 1))
+ result = py_func.wrap_py_func(test_fn, dtypes.int32, (1, 1, 1))
self.assertEqual(3, self.evaluate(result))
result = py_func.wrap_py_func(
- test_fn, dtypes.int64,
+ test_fn, dtypes.int32,
(constant_op.constant(1), 1, constant_op.constant(1)))
self.assertEqual(3, self.evaluate(result))
@@ -53,9 +53,9 @@
return a * b.foo
with self.cached_session() as sess:
- result = py_func.wrap_py_func(test_fn, dtypes.int64, (7, TestClass()))
+ result = py_func.wrap_py_func(test_fn, dtypes.int32, (7, TestClass()))
self.assertEqual(35, self.evaluate(result))
- result = py_func.wrap_py_func(test_fn, dtypes.int64,
+ result = py_func.wrap_py_func(test_fn, dtypes.int32,
(constant_op.constant(7), TestClass()))
self.assertEqual(35, self.evaluate(result))
@@ -70,12 +70,12 @@
return a * b.foo + c * d.foo
with self.cached_session() as sess:
- result = py_func.wrap_py_func(test_fn, dtypes.int64, (7, TestClass(5)), {
+ result = py_func.wrap_py_func(test_fn, dtypes.int32, (7, TestClass(5)), {
'c': 11,
'd': TestClass(13)
})
self.assertEqual(178, self.evaluate(result))
- result = py_func.wrap_py_func(test_fn, dtypes.int64,
+ result = py_func.wrap_py_func(test_fn, dtypes.int32,
(constant_op.constant(7), TestClass(5)), {
'c': constant_op.constant(11),
'd': TestClass(13)
diff --git a/tensorflow/python/client/session_test.py b/tensorflow/python/client/session_test.py
index 347833c..c4a118a 100644
--- a/tensorflow/python/client/session_test.py
+++ b/tensorflow/python/client/session_test.py
@@ -312,6 +312,7 @@
self.assertEqual(None, res[2])
self.assertEqual(44.0, res[1])
+ @test_util.run_v1_only('b/120545219')
def testFetchAttrs(self):
if attr is None:
self.skipTest('attr module is unavailable.')
@@ -340,6 +341,7 @@
self.assertAllEqual(val3, result.field1)
self.assertAllEqual(val2, result.field2)
+ @test_util.run_v1_only('b/120545219')
def testFetchNestedAttrs(self):
if attr is None:
self.skipTest('attr module is unavailable.')
@@ -1024,6 +1026,7 @@
fed_c_val = c.eval(feed_dict={a.name: [[4.0, 4.0]]})
self.assertAllEqual([[16.0, 16.0, 16.0]], fed_c_val)
+ @test_util.run_v1_only('b/120545219')
def testOperationRunMethod(self):
with session.Session():
a = constant_op.constant(1.0, shape=[1, 2])
@@ -1154,6 +1157,7 @@
else:
importer.import_graph_def(gdef, name='import')
+ @test_util.run_v1_only('b/120545219')
def testParallelRunAndSingleBuild(self):
with session.Session() as sess:
c = constant_op.constant(5.0)
@@ -1174,6 +1178,7 @@
for t in threads:
t.join()
+ @test_util.run_v1_only('b/120545219')
def testParallelRunAndParallelBuild(self):
with session.Session() as sess:
c = constant_op.constant(5.0)
@@ -1274,6 +1279,7 @@
with self.assertRaisesRegexp(RuntimeError, 'The Session graph is empty.'):
sess.run({})
+ @test_util.run_v1_only('b/120545219')
def testNotEntered(self):
# pylint: disable=protected-access
self.assertEqual(ops._default_session_stack.get_default(), None)
@@ -1289,6 +1295,7 @@
ValueError, lambda e: 'No default session is registered.' in str(e)):
c_2.eval()
+ @test_util.run_v1_only('b/120545219')
def testInteractive(self):
with ops.device('/cpu:0'):
sess = session.InteractiveSession()
@@ -1301,6 +1308,7 @@
self.assertAllEqual([[24.0]], e.eval())
sess.close()
+ @test_util.run_v1_only('b/120545219')
def testMultipleInteractiveSessionsWarning(self):
# Reinitialize the global state to ensure that the expected warnings will
# be emitted.
@@ -1328,6 +1336,7 @@
sess2.close()
sess.close()
+ @test_util.run_v1_only('b/120545219')
def testInteractivePlacePrunedGraph(self):
sess = session.InteractiveSession()
@@ -1349,6 +1358,7 @@
a.eval()
sess.close()
+ @test_util.run_v1_only('b/120545219')
def testDefaultSessionPlacePrunedGraph(self):
sess = session.Session()
@@ -1769,9 +1779,11 @@
sess.run(a, run_metadata=run_metadata)
self.assertEqual(len(run_metadata.partition_graphs), 0)
+ @test_util.run_v1_only('b/120545219')
def testOutputPartitionGraphsDirect(self):
self.runTestOutputPartitionGraphs(session.Session())
+ @test_util.run_v1_only('b/120545219')
def testOutputPartitionGraphsDistributed(self):
server = server_lib.Server.create_local_server()
self.runTestOutputPartitionGraphs(session.Session(server.target))
@@ -1796,6 +1808,7 @@
del sess1
del sess2
+ @test_util.run_v1_only('b/120545219')
def testAsDefault(self):
c = constant_op.constant(37)
sess = session.Session()
@@ -1821,6 +1834,7 @@
with self.assertRaisesRegexp(TypeError, 'graph must be a tf.Graph'):
session.Session(graph=37)
+ @test_util.run_v1_only('b/120545219')
def testTimeoutWithShortOperations(self):
num_epochs = 5
q = data_flow_ops.FIFOQueue(capacity=50, dtypes=[dtypes.int32], shapes=[()])
@@ -1834,6 +1848,7 @@
sess.run(enqueue_op)
self.assertEqual(sess.run(q.size()), num_epochs * 2)
+ @test_util.run_v1_only('b/120545219')
def testRegisterFetchAndFeedConversionFunctions(self):
class SquaredTensor(object):
@@ -1865,6 +1880,7 @@
squared_eval = sess.partial_run(partial_run, squared_tensor)
self.assertAllClose(np2 * np2, squared_eval)
+ @test_util.run_v1_only('b/120545219')
def testDefaultLogDevicePlacement(self):
class CaptureStderr(str):
@@ -1914,6 +1930,7 @@
self.assertTrue('/job:local/replica:0/task:0/device:CPU:0' in str(log),
str(log))
+ @test_util.run_v1_only('b/120545219')
def testLocalMasterSessionTimeout(self):
# Test that the timeout passed in a config to the session works correctly.
config = config_pb2.ConfigProto(operation_timeout_in_ms=1000)
@@ -1927,6 +1944,7 @@
with self.assertRaises(errors.DeadlineExceededError):
sess.run(dequeued_t)
+ @test_util.run_v1_only('b/120545219')
def testDefaultServerTimeout(self):
# Test that the default server config timeout gets used when no Session
# config is provided.
@@ -1952,9 +1970,11 @@
with self.assertRaisesOpError('has inputs from different frames'):
sess.run(res, feed_dict={data: 1.0})
+ @test_util.run_v1_only('b/120545219')
def testBuildGraphErrorDirect(self):
self.runTestBuildGraphError(session.Session())
+ @test_util.run_v1_only('b/120545219')
def testBuildGraphErrorDist(self):
server = server_lib.Server.create_local_server()
self.runTestBuildGraphError(session.Session(server.target))
@@ -1993,9 +2013,11 @@
result = sess.run(f)
self.assertEqual(result, 2.0)
+ @test_util.run_v1_only('b/120545219')
def testAddFunctionToSession(self):
self.runTestAddFunctionToSession()
+ @test_util.run_v1_only('b/120545219')
def testAddFunctionToGrpcSession(self):
server = server_lib.Server.create_local_server()
self.runTestAddFunctionToSession(server.target)
@@ -2009,6 +2031,7 @@
with session.Session():
pass
+ @test_util.run_v1_only('b/120545219')
def testAutoConvertAndCheckData(self):
with self.cached_session() as sess:
a = array_ops.placeholder(dtype=dtypes.string)
diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 0b6ff30..82ed7da 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -32,7 +32,7 @@
from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2018, 12, 4)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2018, 12, 13)
@tf_export("compat.forward_compatible")
@@ -160,7 +160,7 @@
@tf_export(v1=["disable_v2_behavior"])
def disable_v2_behavior():
- """Enables TensorFlow 2.x behaviors.
+ """Disables TensorFlow 2.x behaviors.
This function can be called at the beginning of the program (before `Tensors`,
`Graphs` or other structures have been created, and before devices have been
diff --git a/tensorflow/python/data/benchmarks/BUILD b/tensorflow/python/data/benchmarks/BUILD
index 5b0500e..fd0eca9 100644
--- a/tensorflow/python/data/benchmarks/BUILD
+++ b/tensorflow/python/data/benchmarks/BUILD
@@ -49,6 +49,20 @@
)
py_test(
+ name = "list_files_benchmark",
+ srcs = ["list_files_benchmark.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:session",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
name = "map_benchmark",
srcs = ["map_benchmark.py"],
srcs_version = "PY2AND3",
diff --git a/tensorflow/python/data/benchmarks/list_files_benchmark.py b/tensorflow/python/data/benchmarks/list_files_benchmark.py
new file mode 100644
index 0000000..0dc2147
--- /dev/null
+++ b/tensorflow/python/data/benchmarks/list_files_benchmark.py
@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.list_files()`."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from os import path
+from os import makedirs
+import shutil
+import time
+import tempfile
+
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+class ListFilesBenchmark(test.Benchmark):
+ """Benchmarks for `tf.data.Dataset.list_files()`."""
+
+ def benchmarkNestedDirectories(self):
+ tmp_dir = tempfile.mkdtemp()
+ width = 1024
+ depth = 16
+ for i in range(width):
+ for j in range(depth):
+ new_base = path.join(tmp_dir, str(i),
+ *[str(dir_name) for dir_name in range(j)])
+ makedirs(new_base)
+ child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
+ for f in child_files:
+ filename = path.join(new_base, f)
+ open(filename, 'w').close()
+ patterns = [
+ path.join(tmp_dir, path.join(*['**'
+ for _ in range(depth)]), suffix)
+ for suffix in ['*.txt', '*.log']
+ ]
+ deltas = []
+ iters = 3
+ for _ in range(iters):
+ with ops.Graph().as_default():
+ dataset = dataset_ops.Dataset.list_files(patterns)
+ next_element = dataset.make_one_shot_iterator().get_next()
+ with session.Session() as sess:
+ sub_deltas = []
+ while True:
+ try:
+ start = time.time()
+ sess.run(next_element)
+ end = time.time()
+ sub_deltas.append(end - start)
+ except errors.OutOfRangeError:
+ break
+ deltas.append(sub_deltas)
+ median_deltas = np.median(deltas, axis=0)
+ print('Nested directory size (width*depth): %d*%d Median wall time: '
+ '%fs (read first filename), %fs (read second filename), avg %fs'
+ ' (read %d more filenames)' %
+ (width, depth, median_deltas[0], median_deltas[1],
+ np.average(median_deltas[2:]), len(median_deltas) - 2))
+ self.report_benchmark(
+ iters=iters,
+ wall_time=np.sum(median_deltas),
+ extras={
+ 'read first file:':
+ median_deltas[0],
+ 'read second file:':
+ median_deltas[1],
+ 'avg time for reading %d more filenames:' %
+ (len(median_deltas) - 2):
+ np.average(median_deltas[2:])
+ },
+ name='nested_directory(%d*%d)' % (width, depth))
+ shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
+if __name__ == '__main__':
+ test.main()
diff --git a/tensorflow/python/data/experimental/__init__.py b/tensorflow/python/data/experimental/__init__.py
index 14dfec3..ffc2e5e 100644
--- a/tensorflow/python/data/experimental/__init__.py
+++ b/tensorflow/python/data/experimental/__init__.py
@@ -25,14 +25,20 @@
@@Counter
@@CheckpointInputPipelineHook
@@CsvDataset
+@@DatasetStructure
+@@NestedStructure
@@OptimizationOptions
@@Optional
+@@OptionalStructure
@@RandomDataset
@@Reducer
+@@SparseTensorStructure
@@SqlDataset
@@StatsAggregator
@@StatsOptions
+@@Structure
@@TFRecordWriter
+@@TensorStructure
@@ThreadingOptions
@@bucket_by_sequence_length
@@ -112,8 +118,14 @@
from tensorflow.python.data.experimental.ops.threading_options import ThreadingOptions
from tensorflow.python.data.experimental.ops.unique import unique
from tensorflow.python.data.experimental.ops.writers import TFRecordWriter
+from tensorflow.python.data.ops.dataset_ops import DatasetStructure
from tensorflow.python.data.ops.iterator_ops import get_next_as_optional
from tensorflow.python.data.ops.optional_ops import Optional
+from tensorflow.python.data.ops.optional_ops import OptionalStructure
+from tensorflow.python.data.util.structure import NestedStructure
+from tensorflow.python.data.util.structure import SparseTensorStructure
+from tensorflow.python.data.util.structure import Structure
+from tensorflow.python.data.util.structure import TensorStructure
# pylint: enable=unused-import
from tensorflow.python.util.all_util import remove_undocumented
diff --git a/tensorflow/python/data/experimental/benchmarks/BUILD b/tensorflow/python/data/experimental/benchmarks/BUILD
index 8175116..651dfd6 100644
--- a/tensorflow/python/data/experimental/benchmarks/BUILD
+++ b/tensorflow/python/data/experimental/benchmarks/BUILD
@@ -59,6 +59,22 @@
)
py_test(
+ name = "map_defun_benchmark",
+ srcs = ["map_defun_benchmark.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:functional_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:tensor_spec",
+ "//tensorflow/python/data/experimental/ops:map_defun",
+ "//tensorflow/python/eager:function",
+ ],
+)
+
+py_test(
name = "map_vectorization_benchmark",
srcs = ["map_vectorization_benchmark.py"],
srcs_version = "PY2AND3",
@@ -109,6 +125,20 @@
)
py_test(
+ name = "rejection_resample_benchmark",
+ srcs = ["rejection_resample_benchmark.py"],
+ srcs_version = "PY2AND3",
+ tags = ["no_pip"],
+ deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/experimental/ops:resampling",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//third_party/py/numpy",
+ "@six_archive//:six",
+ ],
+)
+
+py_test(
name = "unbatch_benchmark",
srcs = ["unbatch_benchmark.py"],
srcs_version = "PY2AND3",
diff --git a/tensorflow/python/data/experimental/benchmarks/autotune_benchmark.py b/tensorflow/python/data/experimental/benchmarks/autotune_benchmark.py
index b48ef95..e713494 100644
--- a/tensorflow/python/data/experimental/benchmarks/autotune_benchmark.py
+++ b/tensorflow/python/data/experimental/benchmarks/autotune_benchmark.py
@@ -134,33 +134,33 @@
a = (np.random.rand(1, 8 * k), np.random.rand(8 * k, 1))
b = (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))
c = (np.random.rand(1, 2 * k), np.random.rand(2 * k, 1))
- dataset = dataset_ops.Dataset.from_tensors((a, b, c)).repeat()
+ dataset_a = dataset_ops.Dataset.from_tensors(a).repeat()
+ dataset_b = dataset_ops.Dataset.from_tensors(b).repeat()
+ dataset_c = dataset_ops.Dataset.from_tensors(c).repeat()
- def f1(a, b, c):
- x, y = a
- return math_ops.matmul(x, y), b, c
+ def f1(x, y):
+ return math_ops.matmul(x, y)
- def f2(a, b, c):
+ def f2(a, b):
x, y = b
- return a, math_ops.matmul(x, y), c
+ return a, math_ops.matmul(x, y)
- def f3(a, b, c):
- x, y = c
- return a, b, math_ops.matmul(x, y)
-
+ dataset = dataset_a
dataset = dataset.map(f1, num_parallel_calls=optimization.AUTOTUNE)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
num_parallel_calls=optimization.AUTOTUNE,
cycle_length=2)
+ dataset = dataset_ops.Dataset.zip((dataset, dataset_b))
dataset = dataset.map(f2, num_parallel_calls=optimization.AUTOTUNE)
dataset = dataset_ops.Dataset.range(1).repeat().interleave(
lambda _: dataset,
num_parallel_calls=optimization.AUTOTUNE,
cycle_length=2)
- dataset = dataset.map(f3, num_parallel_calls=optimization.AUTOTUNE)
+ dataset = dataset_ops.Dataset.zip((dataset, dataset_c))
+ dataset = dataset.map(f2, num_parallel_calls=optimization.AUTOTUNE)
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
diff --git a/tensorflow/python/data/experimental/benchmarks/map_and_batch_benchmark.py b/tensorflow/python/data/experimental/benchmarks/map_and_batch_benchmark.py
index fbd06a5..b17f2bc 100644
--- a/tensorflow/python/data/experimental/benchmarks/map_and_batch_benchmark.py
+++ b/tensorflow/python/data/experimental/benchmarks/map_and_batch_benchmark.py
@@ -26,6 +26,7 @@
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -139,31 +140,11 @@
num_iters = 1024 // (
(element_size * batch_size) // min(num_calls, inter_op))
- dataset = make_base_dataset(element_size)
- chained_dataset = dataset.map(
+ fused_dataset = make_base_dataset(element_size)
+ fused_dataset = fused_dataset.map(
math_ops.matmul,
num_parallel_calls=num_calls).batch(batch_size=batch_size)
- chained_iterator = dataset_ops.make_one_shot_iterator(chained_dataset)
- chained_get_next = chained_iterator.get_next()
- chained_deltas = []
- with session.Session(
- config=config_pb2.ConfigProto(
- inter_op_parallelism_threads=inter_op,
- use_per_session_threads=True)) as sess:
- for _ in range(5):
- sess.run(chained_get_next.op)
- for _ in range(num_iters):
- start = time.time()
- sess.run(chained_get_next.op)
- end = time.time()
- chained_deltas.append(end - start)
-
- fused_dataset = dataset.apply(
- batching.map_and_batch(
- math_ops.matmul,
- num_parallel_calls=num_calls,
- batch_size=batch_size))
fused_iterator = dataset_ops.make_one_shot_iterator(fused_dataset)
fused_get_next = fused_iterator.get_next()
@@ -181,6 +162,28 @@
end = time.time()
fused_deltas.append(end - start)
+ # `map_and_batch_fusion` is optimized by default. To get the chained
+ # dataset, with have to disable it.
+ options = dataset_ops.Options()
+ options.experimental_optimization = OptimizationOptions()
+ options.experimental_optimization.map_and_batch_fusion = False
+ chained_dataset = fused_dataset.with_options(options)
+ chained_iterator = dataset_ops.make_one_shot_iterator(chained_dataset)
+ chained_get_next = chained_iterator.get_next()
+
+ chained_deltas = []
+ with session.Session(
+ config=config_pb2.ConfigProto(
+ inter_op_parallelism_threads=inter_op,
+ use_per_session_threads=True)) as sess:
+ for _ in range(5):
+ sess.run(chained_get_next.op)
+ for _ in range(num_iters):
+ start = time.time()
+ sess.run(chained_get_next.op)
+ end = time.time()
+ chained_deltas.append(end - start)
+
print(
"batch size: %d, num parallel calls: %d, inter-op parallelism: %d, "
"element size: %d, num iters: %d\nchained wall time: %f (median), "
diff --git a/tensorflow/python/data/experimental/benchmarks/map_defun_benchmark.py b/tensorflow/python/data/experimental/benchmarks/map_defun_benchmark.py
new file mode 100644
index 0000000..21e7dda
--- /dev/null
+++ b/tensorflow/python/data/experimental/benchmarks/map_defun_benchmark.py
@@ -0,0 +1,73 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for MapDefunOp."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+
+from tensorflow.python.data.experimental.ops import map_defun
+from tensorflow.python.eager import function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+# TODO(b/119837791): Add eager benchmarks too.
+class MapDefunBenchmark(test.Benchmark):
+ """Benchmarks for MapDefunOp."""
+
+ def _run(self, op, name=None, num_iters=3000):
+ for _ in range(5):
+ self.evaluate(op)
+ start = time.time()
+ for _ in range(num_iters):
+ self.evaluate(op)
+ end = time.time()
+ mean_us = (end - start) * 1e6 / num_iters
+ self.report_benchmark(
+ name=name,
+ iters=num_iters,
+ wall_time=mean_us,
+ extras={"examples_per_sec": num_iters / (end - start)})
+
+ def benchmarkDefunVsMapFn(self):
+ """Benchmarks to compare the performance of MapDefun vs tf.map_fn."""
+
+ @function.defun(input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+ def defun(x):
+ return array_ops.identity(x)
+
+ def map_fn(x):
+ return array_ops.identity(x)
+
+ base = math_ops.range(100)
+ for input_size in [10, 100, 1000, 10000]:
+ num_iters = 100000 // input_size
+ map_defun_op = map_defun.map_defun(defun, [base], [dtypes.int32], [()])
+ map_fn_op = functional_ops.map_fn(map_fn, base)
+
+ self._run(
+ map_defun_op, "with_defun_size_%d" % input_size, num_iters=num_iters)
+ self._run(
+ map_fn_op, "without_defun_size_%d" % input_size, num_iters=num_iters)
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/data/experimental/benchmarks/map_vectorization_benchmark.py b/tensorflow/python/data/experimental/benchmarks/map_vectorization_benchmark.py
index 47ec639..a60ba0a 100644
--- a/tensorflow/python/data/experimental/benchmarks/map_vectorization_benchmark.py
+++ b/tensorflow/python/data/experimental/benchmarks/map_vectorization_benchmark.py
@@ -24,6 +24,7 @@
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import optimization_options
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
from tensorflow.python.framework import constant_op
@@ -115,13 +116,23 @@
def _compare(self, input_dataset, map_fn, batch_size, input_size, str_id):
num_elems = int(np.sum([np.prod(x) for x in input_size]))
name_template = "{}__batch_size_{}_input_element_size_{}_{}"
- unoptimized = input_dataset.map(map_fn).batch(batch_size)
- unoptimized_op = dataset_ops.make_one_shot_iterator(unoptimized).get_next()
- optimized = input_dataset.map(map_fn).batch(batch_size)
+ base_dataset = input_dataset.map(map_fn).batch(batch_size)
+
options = dataset_ops.Options()
- options.experimental_map_vectorization = True
- optimized = optimized.with_options(options)
+ opt_options = optimization_options.OptimizationOptions()
+ # Disable default map_and_batch_fusion optimization
+ opt_options.map_and_batch_fusion = False
+ options.experimental_optimization = opt_options
+ base_dataset = base_dataset.with_options(options)
+
+ unoptimized_op = dataset_ops.make_one_shot_iterator(base_dataset).get_next()
+
+ optimized_options = dataset_ops.Options()
+ opt_options = optimization_options.OptimizationOptions()
+ opt_options.map_vectorization = True
+ optimized_options.experimental_optimization = opt_options
+ optimized = base_dataset.with_options(optimized_options)
optimized_op = dataset_ops.make_one_shot_iterator(optimized).get_next()
unoptimized_time = self._run(
diff --git a/tensorflow/python/data/experimental/benchmarks/optimize_benchmark.py b/tensorflow/python/data/experimental/benchmarks/optimize_benchmark.py
index 2f9b891..5df57a3 100644
--- a/tensorflow/python/data/experimental/benchmarks/optimize_benchmark.py
+++ b/tensorflow/python/data/experimental/benchmarks/optimize_benchmark.py
@@ -28,6 +28,7 @@
from tensorflow.python.platform import test
+# TODO(b/119837791): Add eager benchmarks too.
class OptimizationBenchmark(test.Benchmark):
"""Benchmarks for static optimizations."""
@@ -115,6 +116,46 @@
name="map_and_filter_fusion_{}_chain_length_{}".format(
opt_mark, chain_length))
+ # This benchmark compares the performance of pipeline with multiple chained
+ # filter with and without filter fusion.
+ def benchmarkFilterFusion(self):
+ chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+ for chain_length in chain_lengths:
+ self._benchmarkFilters(chain_length, False)
+ self._benchmarkFilters(chain_length, True)
+
+ def _benchmarkFilterFusion(self, chain_length, optimize_dataset):
+ with ops.Graph().as_default():
+ dataset = dataset_ops.Dataset.from_tensors(5).repeat(None)
+ for _ in range(chain_length):
+ dataset = dataset.filter(lambda x: math_ops.greater_equal(x - 5, 0))
+ if optimize_dataset:
+ options = dataset_ops.Options()
+ options.experimental_filter_fusion = True
+ dataset = dataset.with_options(options)
+
+ iterator = dataset_ops.make_one_shot_iterator(dataset)
+ next_element = iterator.get_next()
+
+ for _ in range(10):
+ self.evaluate(next_element.op)
+ deltas = []
+ for _ in range(100):
+ start = time.time()
+ for _ in range(100):
+ self.evaluate(next_element.op)
+ end = time.time()
+ deltas.append(end - start)
+
+ median_wall_time = np.median(deltas) / 100
+ opt_mark = "opt" if optimize_dataset else "no-opt"
+ print("Filter dataset {} chain length: {} Median wall time: {}".format(
+ opt_mark, chain_length, median_wall_time))
+ self.report_benchmark(
+ iters=1000,
+ wall_time=median_wall_time,
+ name="chain_length_{}_{}".format(opt_mark, chain_length))
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/benchmarks/rejection_resample_benchmark.py b/tensorflow/python/data/experimental/benchmarks/rejection_resample_benchmark.py
new file mode 100644
index 0000000..4cd8c4b
--- /dev/null
+++ b/tensorflow/python/data/experimental/benchmarks/rejection_resample_benchmark.py
@@ -0,0 +1,71 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for `tf.data.experimental.rejection_resample()`."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+
+import numpy as np
+from six.moves import xrange # pylint: disable=redefined-builtin
+
+from tensorflow.python.data.experimental.ops import resampling
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.platform import test
+
+
+def _time_resampling(
+ test_obj, data_np, target_dist, init_dist, num_to_sample):
+ dataset = dataset_ops.Dataset.from_tensor_slices(data_np).repeat()
+
+ # Reshape distribution via rejection sampling.
+ dataset = dataset.apply(
+ resampling.rejection_resample(
+ class_func=lambda x: x,
+ target_dist=target_dist,
+ initial_dist=init_dist,
+ seed=142))
+
+ get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
+
+ with test_obj.test_session() as sess:
+ start_time = time.time()
+ for _ in xrange(num_to_sample):
+ sess.run(get_next)
+ end_time = time.time()
+
+ return end_time - start_time
+
+
+class RejectionResampleBenchmark(test.Benchmark):
+ """Benchmarks for `tf.data.experimental.rejection_resample()`."""
+
+ def benchmarkResamplePerformance(self):
+ init_dist = [0.25, 0.25, 0.25, 0.25]
+ target_dist = [0.0, 0.0, 0.0, 1.0]
+ num_classes = len(init_dist)
+ # We don't need many samples to test a dirac-delta target distribution
+ num_samples = 1000
+ data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+ resample_time = _time_resampling(
+ self, data_np, target_dist, init_dist, num_to_sample=1000)
+
+ self.report_benchmark(iters=1000, wall_time=resample_time, name="resample")
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py b/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py
index c36a325..6f80df5 100644
--- a/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py
+++ b/tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for `tf.data.experimental.unbatch()`."""
+"""Benchmarks for `tf.data.experimental.unbatch()`."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
diff --git a/tensorflow/python/data/experimental/kernel_tests/BUILD b/tensorflow/python/data/experimental/kernel_tests/BUILD
index c76e576..9362a3e 100644
--- a/tensorflow/python/data/experimental/kernel_tests/BUILD
+++ b/tensorflow/python/data/experimental/kernel_tests/BUILD
@@ -1,12 +1,12 @@
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+
package(default_visibility = ["//tensorflow:internal"])
licenses(["notice"]) # Apache 2.0
exports_files(["LICENSE"])
-load("//tensorflow:tensorflow.bzl", "cuda_py_test")
-load("//tensorflow:tensorflow.bzl", "py_test")
-
py_test(
name = "bucket_by_sequence_length_test",
size = "medium",
@@ -130,26 +130,6 @@
)
py_test(
- name = "filter_dataset_op_test",
- size = "medium",
- srcs = ["filter_dataset_op_test.py"],
- srcs_version = "PY2AND3",
- deps = [
- "//tensorflow/python:array_ops",
- "//tensorflow/python:client_testlib",
- "//tensorflow/python:errors",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:io_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:util",
- "//tensorflow/python/data/experimental/ops:optimization",
- "//tensorflow/python/data/kernel_tests:test_base",
- "//tensorflow/python/data/ops:dataset_ops",
- "//third_party/py/numpy",
- ],
-)
-
-py_test(
name = "get_single_element_test",
size = "small",
srcs = ["get_single_element_test.py"],
@@ -545,6 +525,7 @@
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:script_ops",
"//tensorflow/python:sparse_tensor",
"//tensorflow/python/data/experimental/ops:scan_ops",
"//tensorflow/python/data/kernel_tests:test_base",
@@ -621,7 +602,7 @@
py_test(
name = "stats_dataset_ops_test",
- size = "medium",
+ size = "large",
srcs = ["stats_dataset_ops_test.py"],
srcs_version = "PY2AND3",
tags = [
@@ -715,3 +696,14 @@
"//tensorflow/python/data/ops:dataset_ops",
],
)
+
+cuda_py_test(
+ name = "wrap_unwrap_test",
+ size = "small",
+ srcs = ["wrap_unwrap_test.py"],
+ additional_deps = [
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ ],
+)
diff --git a/tensorflow/python/data/experimental/kernel_tests/cardinality_test.py b/tensorflow/python/data/experimental/kernel_tests/cardinality_test.py
index 943f0f1..4a8296d 100644
--- a/tensorflow/python/data/experimental/kernel_tests/cardinality_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/cardinality_test.py
@@ -22,9 +22,11 @@
from tensorflow.python.data.experimental.ops import cardinality
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class NumElementsTest(test_base.DatasetTestBase, parameterized.TestCase):
"""Tests for `tf.data.experimental.cardinality()`."""
diff --git a/tensorflow/python/data/experimental/kernel_tests/copy_to_device_test.py b/tensorflow/python/data/experimental/kernel_tests/copy_to_device_test.py
index b8166fe..786eae1 100644
--- a/tensorflow/python/data/experimental/kernel_tests/copy_to_device_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/copy_to_device_test.py
@@ -33,6 +33,7 @@
from tensorflow.python.util import compat as util_compat
+# TODO(b/119837791): add eager coverage when supported.
class CopyToDeviceTest(test_base.DatasetTestBase):
@test_util.run_deprecated_v1
diff --git a/tensorflow/python/data/experimental/kernel_tests/counter_test.py b/tensorflow/python/data/experimental/kernel_tests/counter_test.py
index 49e1f22..436fa50 100644
--- a/tensorflow/python/data/experimental/kernel_tests/counter_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/counter_test.py
@@ -19,35 +19,31 @@
from tensorflow.python.data.experimental.ops import counter
from tensorflow.python.data.kernel_tests import test_base
-from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class CounterTest(test_base.DatasetTestBase):
- @test_util.run_deprecated_v1
def testCounter(self):
"""Test dataset construction using `count`."""
- iterator = dataset_ops.make_one_shot_iterator(
- counter.Counter(start=3, step=4))
- get_next = iterator.get_next()
- self.assertEqual([], get_next.shape.as_list())
- self.assertEqual(dtypes.int64, get_next.dtype)
+ dataset = counter.Counter(start=3, step=4)
+ self.assertEqual([], dataset.output_shapes.as_list())
+ self.assertEqual(dtypes.int64, dataset.output_types)
+ get_next = self.getNext(dataset)
- negative_iterator = dataset_ops.make_one_shot_iterator(
- counter.Counter(start=0, step=-1))
- negative_get_next = negative_iterator.get_next()
+ negative_dataset = counter.Counter(start=0, step=-1)
+ negative_get_next = self.getNext(negative_dataset)
- with self.cached_session() as sess:
- self.assertEqual(3, self.evaluate(get_next))
- self.assertEqual(3 + 4, self.evaluate(get_next))
- self.assertEqual(3 + 2 * 4, self.evaluate(get_next))
+ self.assertEqual(3, self.evaluate(get_next()))
+ self.assertEqual(3 + 4, self.evaluate(get_next()))
+ self.assertEqual(3 + 2 * 4, self.evaluate(get_next()))
- self.assertEqual(0, self.evaluate(negative_get_next))
- self.assertEqual(-1, self.evaluate(negative_get_next))
- self.assertEqual(-2, self.evaluate(negative_get_next))
+ self.assertEqual(0, self.evaluate(negative_get_next()))
+ self.assertEqual(-1, self.evaluate(negative_get_next()))
+ self.assertEqual(-2, self.evaluate(negative_get_next()))
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py
index 2144342..df69a9d 100644
--- a/tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/directed_interleave_dataset_test.py
@@ -28,9 +28,9 @@
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class DirectedInterleaveDatasetTest(test_base.DatasetTestBase):
- @test_util.run_deprecated_v1
def testBasic(self):
selector_dataset = dataset_ops.Dataset.range(10).repeat(100)
input_datasets = [
@@ -38,16 +38,13 @@
]
dataset = interleave_ops._DirectedInterleaveDataset(selector_dataset,
input_datasets)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
+ next_element = self.getNext(dataset)
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for _ in range(100):
- for i in range(10):
- self.assertEqual(i, self.evaluate(next_element))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for _ in range(100):
+ for i in range(10):
+ self.assertEqual(i, self.evaluate(next_element()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
def _normalize(self, vec):
return vec / vec.sum()
@@ -67,19 +64,16 @@
for i in range(num_datasets)
], weights)
dataset = dataset.take(num_samples)
- iterator = dataset_ops.make_one_shot_iterator(dataset)
- next_element = iterator.get_next()
- with self.cached_session() as sess:
- freqs = np.zeros([num_datasets])
- for _ in range(num_samples):
- freqs[self.evaluate(next_element)] += 1
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ next_element = self.getNext(dataset)
+ freqs = np.zeros([num_datasets])
+ for _ in range(num_samples):
+ freqs[self.evaluate(next_element())] += 1
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
return freqs
- @test_util.run_deprecated_v1
def testSampleFromDatasets(self):
random_seed.set_random_seed(1619)
num_samples = 5000
@@ -99,21 +93,17 @@
freqs = self._testSampleFromDatasetsHelper(probs_ds, classes, num_samples)
self.assertLess(self._chi2(probs, freqs / num_samples), 1e-2)
- @test_util.run_deprecated_v1
def testSelectFromDatasets(self):
words = [b"foo", b"bar", b"baz"]
datasets = [dataset_ops.Dataset.from_tensors(w).repeat() for w in words]
choice_array = np.random.randint(3, size=(15,), dtype=np.int64)
choice_dataset = dataset_ops.Dataset.from_tensor_slices(choice_array)
dataset = interleave_ops.choose_from_datasets(datasets, choice_dataset)
- iterator = dataset_ops.make_one_shot_iterator(dataset)
- next_element = iterator.get_next()
-
- with self.cached_session() as sess:
- for i in choice_array:
- self.assertEqual(words[i], self.evaluate(next_element))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ next_element = self.getNext(dataset)
+ for i in choice_array:
+ self.assertEqual(words[i], self.evaluate(next_element()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
def testErrors(self):
with self.assertRaisesRegexp(ValueError,
diff --git a/tensorflow/python/data/experimental/kernel_tests/enumerate_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/enumerate_dataset_test.py
index 2574209..cbc048e 100644
--- a/tensorflow/python/data/experimental/kernel_tests/enumerate_dataset_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/enumerate_dataset_test.py
@@ -22,37 +22,28 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class EnumerateDatasetTest(test_base.DatasetTestBase):
- @test_util.run_deprecated_v1
def testEnumerateDataset(self):
components = (["a", "b"], [1, 2], [37.0, 38])
start = constant_op.constant(20, dtype=dtypes.int64)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(components).apply(
- enumerate_ops.enumerate_dataset(start)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensor_slices(components).apply(
+ enumerate_ops.enumerate_dataset(start))
- self.assertEqual(dtypes.int64, get_next[0].dtype)
- self.assertEqual((), get_next[0].shape)
+ self.assertEqual(dtypes.int64, dataset.output_types[0])
+ self.assertEqual((), dataset.output_shapes[0])
self.assertEqual([tensor_shape.TensorShape([])] * 3,
- [t.shape for t in get_next[1]])
+ [shape for shape in dataset.output_shapes[1]])
- with self.cached_session() as sess:
- self.evaluate(init_op)
- self.assertEqual((20, (b"a", 1, 37.0)), self.evaluate(get_next))
- self.assertEqual((21, (b"b", 2, 38.0)), self.evaluate(get_next))
-
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ self.assertDatasetProduces(dataset, [(20, (b"a", 1, 37.0)),
+ (21, (b"b", 2, 38.0))])
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/filter_dataset_op_test.py b/tensorflow/python/data/experimental/kernel_tests/filter_dataset_op_test.py
deleted file mode 100644
index 357b5f1..0000000
--- a/tensorflow/python/data/experimental/kernel_tests/filter_dataset_op_test.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# 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.
-# ==============================================================================
-"""Benchmarks FilterDataset input pipeline op."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import time
-
-import numpy as np
-
-from tensorflow.python.client import session
-from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.platform import test
-
-
-class FilterBenchmark(test.Benchmark):
-
- # This benchmark compares the performance of pipeline with multiple chained
- # filter with and without filter fusion.
- def benchmarkFilters(self):
- chain_lengths = [0, 1, 2, 5, 10, 20, 50]
- for chain_length in chain_lengths:
- self._benchmarkFilters(chain_length, False)
- self._benchmarkFilters(chain_length, True)
-
- def _benchmarkFilters(self, chain_length, optimize_dataset):
- with ops.Graph().as_default():
- dataset = dataset_ops.Dataset.from_tensors(5).repeat(None)
- for _ in range(chain_length):
- dataset = dataset.filter(lambda x: math_ops.greater_equal(x - 5, 0))
- if optimize_dataset:
- dataset = dataset.apply(optimization.optimize(["filter_fusion"]))
-
- iterator = dataset_ops.make_one_shot_iterator(dataset)
- next_element = iterator.get_next()
-
- with session.Session() as sess:
- for _ in range(10):
- self.evaluate(next_element.op)
- deltas = []
- for _ in range(100):
- start = time.time()
- for _ in range(100):
- self.evaluate(next_element.op)
- end = time.time()
- deltas.append(end - start)
-
- median_wall_time = np.median(deltas) / 100
- opt_mark = "opt" if optimize_dataset else "no-opt"
- print("Filter dataset {} chain length: {} Median wall time: {}".format(
- opt_mark, chain_length, median_wall_time))
- self.report_benchmark(
- iters=1000,
- wall_time=median_wall_time,
- name="benchmark_filter_dataset_chain_latency_{}_{}".format(
- opt_mark, chain_length))
-
-
-if __name__ == "__main__":
- test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py b/tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py
index ef57656..3e2cf77 100644
--- a/tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/get_single_element_test.py
@@ -22,7 +22,6 @@
from tensorflow.python.data.experimental.ops import get_single_element
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
@@ -30,6 +29,7 @@
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class GetSingleElementTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
@@ -40,34 +40,25 @@
("MoreThanOne", 0, 2, errors.InvalidArgumentError,
"Dataset had more than one element."),
)
- @test_util.run_deprecated_v1
def testGetSingleElement(self, skip, take, error=None, error_msg=None):
- skip_t = array_ops.placeholder(dtypes.int64, shape=[])
- take_t = array_ops.placeholder(dtypes.int64, shape=[])
def make_sparse(x):
x_1d = array_ops.reshape(x, [1])
x_2d = array_ops.reshape(x, [1, 1])
return sparse_tensor.SparseTensor(x_2d, x_1d, x_1d)
- dataset = dataset_ops.Dataset.range(100).skip(skip_t).map(
- lambda x: (x * x, make_sparse(x))).take(take_t)
- element = get_single_element.get_single_element(dataset)
-
- with self.cached_session() as sess:
- if error is None:
- dense_val, sparse_val = sess.run(
- element, feed_dict={
- skip_t: skip,
- take_t: take
- })
- self.assertEqual(skip * skip, dense_val)
- self.assertAllEqual([[skip]], sparse_val.indices)
- self.assertAllEqual([skip], sparse_val.values)
- self.assertAllEqual([skip], sparse_val.dense_shape)
- else:
- with self.assertRaisesRegexp(error, error_msg):
- sess.run(element, feed_dict={skip_t: skip, take_t: take})
+ dataset = dataset_ops.Dataset.range(100).skip(
+ skip).map(lambda x: (x * x, make_sparse(x))).take(take)
+ if error is None:
+ dense_val, sparse_val = self.evaluate(
+ get_single_element.get_single_element(dataset))
+ self.assertEqual(skip * skip, dense_val)
+ self.assertAllEqual([[skip]], sparse_val.indices)
+ self.assertAllEqual([skip], sparse_val.values)
+ self.assertAllEqual([skip], sparse_val.dense_shape)
+ else:
+ with self.assertRaisesRegexp(error, error_msg):
+ self.evaluate(get_single_element.get_single_element(dataset))
def testWindow(self):
"""Test that `get_single_element()` can consume a nested dataset."""
diff --git a/tensorflow/python/data/experimental/kernel_tests/group_by_reducer_test.py b/tensorflow/python/data/experimental/kernel_tests/group_by_reducer_test.py
index 8507df3..4194f06 100644
--- a/tensorflow/python/data/experimental/kernel_tests/group_by_reducer_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/group_by_reducer_test.py
@@ -33,19 +33,9 @@
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class GroupByReducerTest(test_base.DatasetTestBase):
- def checkResults(self, dataset, shapes, values):
- self.assertEqual(shapes, dataset.output_shapes)
- get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
- with self.cached_session() as sess:
- for expected in values:
- got = self.evaluate(get_next)
- self.assertEqual(got, expected)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
-
- @test_util.run_deprecated_v1
def testSum(self):
reducer = grouping.Reducer(
init_func=lambda _: np.int64(0),
@@ -54,10 +44,11 @@
for i in range(1, 11):
dataset = dataset_ops.Dataset.range(2 * i).apply(
grouping.group_by_reducer(lambda x: x % 2, reducer))
- self.checkResults(
- dataset, shapes=tensor_shape.scalar(), values=[(i - 1) * i, i * i])
+ self.assertDatasetProduces(
+ dataset,
+ expected_shapes=tensor_shape.scalar(),
+ expected_output=[(i - 1) * i, i * i])
- @test_util.run_deprecated_v1
def testAverage(self):
def reduce_fn(x, y):
@@ -72,10 +63,11 @@
dataset = dataset_ops.Dataset.range(2 * i).apply(
grouping.group_by_reducer(
lambda x: math_ops.cast(x, dtypes.int64) % 2, reducer))
- self.checkResults(
- dataset, shapes=tensor_shape.scalar(), values=[i - 1, i])
+ self.assertDatasetProduces(
+ dataset,
+ expected_shapes=tensor_shape.scalar(),
+ expected_output=[i - 1, i])
- @test_util.run_deprecated_v1
def testConcat(self):
components = np.array(list("abcdefghijklmnopqrst")).view(np.chararray)
reducer = grouping.Reducer(
@@ -87,12 +79,11 @@
(dataset_ops.Dataset.from_tensor_slices(components),
dataset_ops.Dataset.range(2 * i))).apply(
grouping.group_by_reducer(lambda x, y: y % 2, reducer))
- self.checkResults(
+ self.assertDatasetProduces(
dataset,
- shapes=tensor_shape.scalar(),
- values=[b"acegikmoqs" [:i], b"bdfhjlnprt" [:i]])
+ expected_shapes=tensor_shape.scalar(),
+ expected_output=[b"acegikmoqs" [:i], b"bdfhjlnprt" [:i]])
- @test_util.run_deprecated_v1
def testSparseSum(self):
def _sparse(i):
return sparse_tensor.SparseTensorValue(
@@ -107,10 +98,11 @@
for i in range(1, 11):
dataset = dataset_ops.Dataset.range(2 * i).map(_sparse).apply(
grouping.group_by_reducer(lambda x: x.values[0] % 2, reducer))
- self.checkResults(
- dataset, shapes=tensor_shape.scalar(), values=[(i - 1) * i, i * i])
+ self.assertDatasetProduces(
+ dataset,
+ expected_shapes=tensor_shape.scalar(),
+ expected_output=[(i - 1) * i, i * i])
- @test_util.run_deprecated_v1
def testChangingStateShape(self):
def reduce_fn(x, _):
@@ -130,14 +122,12 @@
grouping.group_by_reducer(lambda x: x, reducer))
self.assertEqual([None], dataset.output_shapes[0].as_list())
self.assertIs(None, dataset.output_shapes[1].ndims)
- iterator = dataset_ops.make_one_shot_iterator(dataset)
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- x, y = self.evaluate(get_next)
- self.assertAllEqual([0] * (2**i), x)
- self.assertAllEqual(np.array(1, ndmin=i), y)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ get_next = self.getNext(dataset)
+ x, y = self.evaluate(get_next())
+ self.assertAllEqual([0] * (2**i), x)
+ self.assertAllEqual(np.array(1, ndmin=i), y)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testTypeMismatch(self):
reducer = grouping.Reducer(
@@ -194,11 +184,10 @@
dataset = dataset_ops.Dataset.zip(
(dataset_ops.Dataset.range(10), dataset_ops.Dataset.range(10))).apply(
grouping.group_by_reducer(lambda x, y: np.int64(0), reducer))
- get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
- with self.cached_session() as sess:
- x, y = self.evaluate(get_next)
- self.assertAllEqual(x, np.asarray([x for x in range(10)]))
- self.assertEqual(y, 45)
+ get_next = self.getNext(dataset)
+ x, y = self.evaluate(get_next())
+ self.assertAllEqual(x, np.asarray([x for x in range(10)]))
+ self.assertEqual(y, 45)
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/group_by_window_test.py b/tensorflow/python/data/experimental/kernel_tests/group_by_window_test.py
index cbb79e5..d127070 100644
--- a/tensorflow/python/data/experimental/kernel_tests/group_by_window_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/group_by_window_test.py
@@ -37,6 +37,7 @@
# NOTE(mrry): These tests are based on the tests in bucket_ops_test.py.
# Currently, they use a constant batch size, though should be made to use a
# different batch size per key.
+@test_util.run_all_in_graph_and_eager_modes
class GroupByWindowTest(test_base.DatasetTestBase):
def _dynamicPad(self, bucket, window, window_size):
@@ -50,101 +51,87 @@
32, (tensor_shape.TensorShape([]), tensor_shape.TensorShape(
[None]), tensor_shape.TensorShape([3])))))
- @test_util.run_deprecated_v1
def testSingleBucket(self):
def _map_fn(v):
return (v, array_ops.fill([v], v),
array_ops.fill([3], string_ops.as_string(v)))
- input_dataset = (
- dataset_ops.Dataset.from_tensor_slices(math_ops.range(32)).map(_map_fn))
+ input_dataset = dataset_ops.Dataset.from_tensor_slices(
+ math_ops.range(32)).map(_map_fn)
bucketed_dataset = input_dataset.apply(
grouping.group_by_window(
lambda x, y, z: 0,
lambda k, bucket: self._dynamicPad(k, bucket, 32), 32))
+ get_next = self.getNext(bucketed_dataset)
- iterator = dataset_ops.make_initializable_iterator(bucketed_dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ which_bucket, bucketed_values = self.evaluate(get_next())
- with self.cached_session() as sess:
- self.evaluate(init_op)
+ self.assertEqual(0, which_bucket)
- which_bucket, bucketed_values = self.evaluate(get_next)
+ expected_scalar_int = np.arange(32, dtype=np.int64)
+ expected_unk_int64 = np.zeros((32, 31)).astype(np.int64)
+ for i in range(32):
+ expected_unk_int64[i, :i] = i
+ expected_vec3_str = np.vstack(3 * [np.arange(32).astype(bytes)]).T
- self.assertEqual(0, which_bucket)
+ self.assertAllEqual(expected_scalar_int, bucketed_values[0])
+ self.assertAllEqual(expected_unk_int64, bucketed_values[1])
+ self.assertAllEqual(expected_vec3_str, bucketed_values[2])
- expected_scalar_int = np.arange(32, dtype=np.int64)
- expected_unk_int64 = np.zeros((32, 31)).astype(np.int64)
- for i in range(32):
- expected_unk_int64[i, :i] = i
- expected_vec3_str = np.vstack(3 * [np.arange(32).astype(bytes)]).T
-
- self.assertAllEqual(expected_scalar_int, bucketed_values[0])
- self.assertAllEqual(expected_unk_int64, bucketed_values[1])
- self.assertAllEqual(expected_vec3_str, bucketed_values[2])
-
- @test_util.run_deprecated_v1
def testEvenOddBuckets(self):
def _map_fn(v):
return (v, array_ops.fill([v], v),
array_ops.fill([3], string_ops.as_string(v)))
- input_dataset = (
- dataset_ops.Dataset.from_tensor_slices(math_ops.range(64)).map(_map_fn))
+ input_dataset = dataset_ops.Dataset.from_tensor_slices(
+ math_ops.range(64)).map(_map_fn)
bucketed_dataset = input_dataset.apply(
grouping.group_by_window(
lambda x, y, z: math_ops.cast(x % 2, dtypes.int64),
lambda k, bucket: self._dynamicPad(k, bucket, 32), 32))
- iterator = dataset_ops.make_initializable_iterator(bucketed_dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(bucketed_dataset)
- with self.cached_session() as sess:
- self.evaluate(init_op)
+ # Get two minibatches (one containing even values, one containing odds)
+ which_bucket_even, bucketed_values_even = self.evaluate(get_next())
+ which_bucket_odd, bucketed_values_odd = self.evaluate(get_next())
- # Get two minibatches (one containing even values, one containing odds)
- which_bucket_even, bucketed_values_even = self.evaluate(get_next)
- which_bucket_odd, bucketed_values_odd = self.evaluate(get_next)
+ # Count number of bucket_tensors.
+ self.assertEqual(3, len(bucketed_values_even))
+ self.assertEqual(3, len(bucketed_values_odd))
- # Count number of bucket_tensors.
- self.assertEqual(3, len(bucketed_values_even))
- self.assertEqual(3, len(bucketed_values_odd))
+ # Ensure bucket 0 was used for all minibatch entries.
+ self.assertAllEqual(0, which_bucket_even)
+ self.assertAllEqual(1, which_bucket_odd)
- # Ensure bucket 0 was used for all minibatch entries.
- self.assertAllEqual(0, which_bucket_even)
- self.assertAllEqual(1, which_bucket_odd)
+ # Test the first bucket outputted, the events starting at 0
+ expected_scalar_int = np.arange(0, 32 * 2, 2, dtype=np.int64)
+ expected_unk_int64 = np.zeros((32, 31 * 2)).astype(np.int64)
+ for i in range(0, 32):
+ expected_unk_int64[i, :2 * i] = 2 * i
+ expected_vec3_str = np.vstack(
+ 3 * [np.arange(0, 32 * 2, 2).astype(bytes)]).T
- # Test the first bucket outputted, the events starting at 0
- expected_scalar_int = np.arange(0, 32 * 2, 2, dtype=np.int64)
- expected_unk_int64 = np.zeros((32, 31 * 2)).astype(np.int64)
- for i in range(0, 32):
- expected_unk_int64[i, :2 * i] = 2 * i
- expected_vec3_str = np.vstack(
- 3 * [np.arange(0, 32 * 2, 2).astype(bytes)]).T
+ self.assertAllEqual(expected_scalar_int, bucketed_values_even[0])
+ self.assertAllEqual(expected_unk_int64, bucketed_values_even[1])
+ self.assertAllEqual(expected_vec3_str, bucketed_values_even[2])
- self.assertAllEqual(expected_scalar_int, bucketed_values_even[0])
- self.assertAllEqual(expected_unk_int64, bucketed_values_even[1])
- self.assertAllEqual(expected_vec3_str, bucketed_values_even[2])
+ # Test the second bucket outputted, the odds starting at 1
+ expected_scalar_int = np.arange(1, 32 * 2 + 1, 2, dtype=np.int64)
+ expected_unk_int64 = np.zeros((32, 31 * 2 + 1)).astype(np.int64)
+ for i in range(0, 32):
+ expected_unk_int64[i, :2 * i + 1] = 2 * i + 1
+ expected_vec3_str = np.vstack(
+ 3 * [np.arange(1, 32 * 2 + 1, 2).astype(bytes)]).T
- # Test the second bucket outputted, the odds starting at 1
- expected_scalar_int = np.arange(1, 32 * 2 + 1, 2, dtype=np.int64)
- expected_unk_int64 = np.zeros((32, 31 * 2 + 1)).astype(np.int64)
- for i in range(0, 32):
- expected_unk_int64[i, :2 * i + 1] = 2 * i + 1
- expected_vec3_str = np.vstack(
- 3 * [np.arange(1, 32 * 2 + 1, 2).astype(bytes)]).T
+ self.assertAllEqual(expected_scalar_int, bucketed_values_odd[0])
+ self.assertAllEqual(expected_unk_int64, bucketed_values_odd[1])
+ self.assertAllEqual(expected_vec3_str, bucketed_values_odd[2])
- self.assertAllEqual(expected_scalar_int, bucketed_values_odd[0])
- self.assertAllEqual(expected_unk_int64, bucketed_values_odd[1])
- self.assertAllEqual(expected_vec3_str, bucketed_values_odd[2])
-
- @test_util.run_deprecated_v1
def testEvenOddBucketsFilterOutAllOdd(self):
def _map_fn(v):
@@ -164,35 +151,28 @@
"z": tensor_shape.TensorShape([3])
})))
- input_dataset = (
- dataset_ops.Dataset.from_tensor_slices(math_ops.range(128)).map(_map_fn)
- .filter(lambda d: math_ops.equal(d["x"] % 2, 0)))
+ input_dataset = dataset_ops.Dataset.from_tensor_slices(math_ops.range(
+ 128)).map(_map_fn).filter(lambda d: math_ops.equal(d["x"] % 2, 0))
bucketed_dataset = input_dataset.apply(
grouping.group_by_window(
lambda d: math_ops.cast(d["x"] % 2, dtypes.int64),
lambda k, bucket: _dynamic_pad_fn(k, bucket, 32), 32))
- iterator = dataset_ops.make_initializable_iterator(bucketed_dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(bucketed_dataset)
- with self.cached_session() as sess:
- self.evaluate(init_op)
+ # Get two minibatches ([0, 2, ...] and [64, 66, ...])
+ which_bucket0, bucketed_values_even0 = self.evaluate(get_next())
+ which_bucket1, bucketed_values_even1 = self.evaluate(get_next())
- # Get two minibatches ([0, 2, ...] and [64, 66, ...])
- which_bucket0, bucketed_values_even0 = self.evaluate(get_next)
- which_bucket1, bucketed_values_even1 = self.evaluate(get_next)
+ # Ensure that bucket 1 was completely filtered out
+ self.assertAllEqual(0, which_bucket0)
+ self.assertAllEqual(0, which_bucket1)
+ self.assertAllEqual(
+ np.arange(0, 64, 2, dtype=np.int64), bucketed_values_even0["x"])
+ self.assertAllEqual(
+ np.arange(64, 128, 2, dtype=np.int64), bucketed_values_even1["x"])
- # Ensure that bucket 1 was completely filtered out
- self.assertAllEqual(0, which_bucket0)
- self.assertAllEqual(0, which_bucket1)
- self.assertAllEqual(
- np.arange(0, 64, 2, dtype=np.int64), bucketed_values_even0["x"])
- self.assertAllEqual(
- np.arange(64, 128, 2, dtype=np.int64), bucketed_values_even1["x"])
-
- @test_util.run_deprecated_v1
def testDynamicWindowSize(self):
components = np.arange(100).astype(np.int64)
@@ -207,111 +187,81 @@
dataset = dataset_ops.Dataset.from_tensor_slices(components).apply(
grouping.group_by_window(lambda x: x % 2, lambda _, xs: xs.batch(20),
None, window_size_func))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- self.evaluate(init_op)
- with self.assertRaises(errors.OutOfRangeError):
- batches = 0
- while True:
- result = self.evaluate(get_next)
- is_even = all(x % 2 == 0 for x in result)
- is_odd = all(x % 2 == 1 for x in result)
- self.assertTrue(is_even or is_odd)
- expected_batch_size = 5 if is_even else 10
- self.assertEqual(expected_batch_size, result.shape[0])
- batches += 1
+ get_next = self.getNext(dataset)
+ with self.assertRaises(errors.OutOfRangeError):
+ batches = 0
+ while True:
+ result = self.evaluate(get_next())
+ is_even = all(x % 2 == 0 for x in result)
+ is_odd = all(x % 2 == 1 for x in result)
+ self.assertTrue(is_even or is_odd)
+ expected_batch_size = 5 if is_even else 10
+ self.assertEqual(expected_batch_size, result.shape[0])
+ batches += 1
- self.assertEqual(batches, 15)
+ self.assertEqual(batches, 15)
- @test_util.run_deprecated_v1
def testSimple(self):
components = np.random.randint(100, size=(200,)).astype(np.int64)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(components).map(lambda x: x * x)
- .apply(
+ dataset = dataset_ops.Dataset.from_tensor_slices(
+ components).map(lambda x: x * x).apply(
grouping.group_by_window(lambda x: x % 2, lambda _, xs: xs.batch(4),
- 4)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ 4))
+ get_next = self.getNext(dataset)
+ counts = []
+ with self.assertRaises(errors.OutOfRangeError):
+ while True:
+ result = self.evaluate(get_next())
+ self.assertTrue(
+ all(x % 2 == 0 for x in result) or all(x % 2 == 1) for x in result)
+ counts.append(result.shape[0])
- with self.cached_session() as sess:
- self.evaluate(init_op)
- counts = []
- with self.assertRaises(errors.OutOfRangeError):
- while True:
- result = self.evaluate(get_next)
- self.assertTrue(
- all(x % 2 == 0
- for x in result) or all(x % 2 == 1)
- for x in result)
- counts.append(result.shape[0])
+ self.assertEqual(len(components), sum(counts))
+ num_full_batches = len([c for c in counts if c == 4])
+ self.assertGreaterEqual(num_full_batches, 24)
+ self.assertTrue(all(c == 4 for c in counts[:num_full_batches]))
- self.assertEqual(len(components), sum(counts))
- num_full_batches = len([c for c in counts if c == 4])
- self.assertGreaterEqual(num_full_batches, 24)
- self.assertTrue(all(c == 4 for c in counts[:num_full_batches]))
-
- @test_util.run_deprecated_v1
def testImmediateOutput(self):
components = np.array(
[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 2, 2, 0, 0], dtype=np.int64)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(components).repeat(-1).apply(
+ dataset = dataset_ops.Dataset.from_tensor_slices(components).repeat(
+ -1).apply(
grouping.group_by_window(lambda x: x % 3, lambda _, xs: xs.batch(4),
- 4)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ 4))
+ get_next = self.getNext(dataset)
+ # The input is infinite, so this test demonstrates that:
+ # 1. We produce output without having to consume the entire input,
+ # 2. Different buckets can produce output at different rates, and
+ # 3. For deterministic input, the output is deterministic.
+ for _ in range(3):
+ self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next()))
+ self.assertAllEqual([1, 1, 1, 1], self.evaluate(get_next()))
+ self.assertAllEqual([2, 2, 2, 2], self.evaluate(get_next()))
+ self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next()))
- with self.cached_session() as sess:
- self.evaluate(init_op)
- # The input is infinite, so this test demonstrates that:
- # 1. We produce output without having to consume the entire input,
- # 2. Different buckets can produce output at different rates, and
- # 3. For deterministic input, the output is deterministic.
- for _ in range(3):
- self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next))
- self.assertAllEqual([1, 1, 1, 1], self.evaluate(get_next))
- self.assertAllEqual([2, 2, 2, 2], self.evaluate(get_next))
- self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next))
-
- @test_util.run_deprecated_v1
def testSmallGroups(self):
components = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0], dtype=np.int64)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(components).apply(
- grouping.group_by_window(lambda x: x % 2, lambda _, xs: xs.batch(4),
- 4)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensor_slices(components).apply(
+ grouping.group_by_window(lambda x: x % 2, lambda _, xs: xs.batch(4), 4))
+ get_next = self.getNext(dataset)
+ self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next()))
+ self.assertAllEqual([1, 1, 1, 1], self.evaluate(get_next()))
+ # The small outputs at the end are deterministically produced in key
+ # order.
+ self.assertAllEqual([0, 0, 0], self.evaluate(get_next()))
+ self.assertAllEqual([1], self.evaluate(get_next()))
- with self.cached_session() as sess:
- self.evaluate(init_op)
- self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next))
- self.assertAllEqual([1, 1, 1, 1], self.evaluate(get_next))
- # The small outputs at the end are deterministically produced in key
- # order.
- self.assertAllEqual([0, 0, 0], self.evaluate(get_next))
- self.assertAllEqual([1], self.evaluate(get_next))
-
- @test_util.run_deprecated_v1
def testEmpty(self):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(4).apply(
- grouping.group_by_window(lambda _: 0, lambda _, xs: xs, 0)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.range(4).apply(
+ grouping.group_by_window(lambda _: 0, lambda _, xs: xs, 0))
- with self.cached_session() as sess:
- self.evaluate(init_op)
- with self.assertRaisesRegexp(
- errors.InvalidArgumentError,
- "Window size must be greater than zero, but got 0."):
- print(self.evaluate(get_next))
+ get_next = self.getNext(dataset)
+ with self.assertRaisesRegexp(
+ errors.InvalidArgumentError,
+ "Window size must be greater than zero, but got 0."):
+ print(self.evaluate(get_next()))
- @test_util.run_deprecated_v1
def testReduceFuncError(self):
components = np.random.randint(100, size=(200,)).astype(np.int64)
@@ -323,19 +273,13 @@
padded_shapes=(tensor_shape.TensorShape([]),
constant_op.constant([5], dtype=dtypes.int64) * -1))
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(components)
- .map(lambda x: (x, ops.convert_to_tensor([x * x]))).apply(
- grouping.group_by_window(lambda x, _: x % 2, reduce_func, 32)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensor_slices(
+ components).map(lambda x: (x, ops.convert_to_tensor([x * x]))).apply(
+ grouping.group_by_window(lambda x, _: x % 2, reduce_func, 32))
+ get_next = self.getNext(dataset)
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(get_next())
- with self.cached_session() as sess:
- self.evaluate(init_op)
- with self.assertRaises(errors.InvalidArgumentError):
- self.evaluate(get_next)
-
- @test_util.run_deprecated_v1
def testConsumeWindowDatasetMoreThanOnce(self):
components = np.random.randint(50, size=(200,)).astype(np.int64)
@@ -349,26 +293,23 @@
4, padded_shapes=ops.convert_to_tensor([(key + 1) * 10])),
))
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(components)
- .map(lambda x: array_ops.fill([math_ops.cast(x, dtypes.int32)], x))
- .apply(grouping.group_by_window(
+ dataset = dataset_ops.Dataset.from_tensor_slices(
+ components
+ ).map(lambda x: array_ops.fill([math_ops.cast(x, dtypes.int32)], x)).apply(
+ grouping.group_by_window(
lambda x: math_ops.cast(array_ops.shape(x)[0] // 10, dtypes.int64),
- reduce_func, 4)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ reduce_func, 4))
- with self.cached_session() as sess:
- self.evaluate(init_op)
- counts = []
- with self.assertRaises(errors.OutOfRangeError):
- while True:
- tight_result, multiple_of_10_result = self.evaluate(get_next)
- self.assertEqual(0, multiple_of_10_result.shape[1] % 10)
- self.assertAllEqual(tight_result,
- multiple_of_10_result[:, :tight_result.shape[1]])
- counts.append(tight_result.shape[0])
- self.assertEqual(len(components), sum(counts))
+ get_next = self.getNext(dataset)
+ counts = []
+ with self.assertRaises(errors.OutOfRangeError):
+ while True:
+ tight_result, multiple_of_10_result = self.evaluate(get_next())
+ self.assertEqual(0, multiple_of_10_result.shape[1] % 10)
+ self.assertAllEqual(tight_result,
+ multiple_of_10_result[:, :tight_result.shape[1]])
+ counts.append(tight_result.shape[0])
+ self.assertEqual(len(components), sum(counts))
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/ignore_errors_test.py b/tensorflow/python/data/experimental/kernel_tests/ignore_errors_test.py
index 81f580f..1d02f4f 100644
--- a/tensorflow/python/data/experimental/kernel_tests/ignore_errors_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/ignore_errors_test.py
@@ -34,9 +34,9 @@
_NUMPY_RANDOM_SEED = 42
+@test_util.run_all_in_graph_and_eager_modes
class IgnoreErrorsTest(test_base.DatasetTestBase):
- @test_util.run_deprecated_v1
def testMapIgnoreError(self):
components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
@@ -44,18 +44,13 @@
dataset_ops.Dataset.from_tensor_slices(components)
.map(lambda x: array_ops.check_numerics(x, "message")).apply(
error_ops.ignore_errors()))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- self.evaluate(init_op)
- for x in [1., 2., 3., 5.]:
- self.assertEqual(x, self.evaluate(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ for x in [1., 2., 3., 5.]:
+ self.assertEqual(x, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- @test_util.run_deprecated_v1
def testParallelMapIgnoreError(self):
components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
@@ -63,18 +58,13 @@
dataset_ops.Dataset.from_tensor_slices(components).map(
lambda x: array_ops.check_numerics(x, "message"),
num_parallel_calls=2).prefetch(2).apply(error_ops.ignore_errors()))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- self.evaluate(init_op)
- for x in [1., 2., 3., 5.]:
- self.assertEqual(x, self.evaluate(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ for x in [1., 2., 3., 5.]:
+ self.assertEqual(x, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- @test_util.run_deprecated_v1
def testReadFileIgnoreError(self):
def write_string_to_file(value, filename):
@@ -91,28 +81,24 @@
dataset_ops.Dataset.from_tensor_slices(filenames).map(
io_ops.read_file,
num_parallel_calls=2).prefetch(2).apply(error_ops.ignore_errors()))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- # All of the files are present.
- self.evaluate(init_op)
- for filename in filenames:
- self.assertEqual(compat.as_bytes(filename), self.evaluate(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ # All of the files are present.
+ for filename in filenames:
+ self.assertEqual(compat.as_bytes(filename), self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- # Delete one of the files.
- os.remove(filenames[0])
+ # Delete one of the files.
+ os.remove(filenames[0])
- # Attempting to read filenames[0] will fail, but ignore_errors()
- # will catch the error.
- self.evaluate(init_op)
- for filename in filenames[1:]:
- self.assertEqual(compat.as_bytes(filename), self.evaluate(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ # Attempting to read filenames[0] will fail, but ignore_errors()
+ # will catch the error.
+ get_next = self.getNext(dataset)
+ for filename in filenames[1:]:
+ self.assertEqual(compat.as_bytes(filename), self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/make_csv_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/make_csv_dataset_test.py
index e80acce..3b7b335 100644
--- a/tensorflow/python/data/experimental/kernel_tests/make_csv_dataset_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/make_csv_dataset_test.py
@@ -25,16 +25,15 @@
from tensorflow.python.data.experimental.ops import readers
from tensorflow.python.data.kernel_tests import test_base
-from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
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 test_util
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class MakeCsvDatasetTest(test_base.DatasetTestBase):
def _make_csv_dataset(self, filenames, batch_size, num_epochs=1, **kwargs):
@@ -76,7 +75,6 @@
def _verify_output(
self,
- sess,
dataset,
batch_size,
num_epochs,
@@ -84,7 +82,7 @@
expected_output,
expected_keys,
):
- nxt = dataset_ops.make_one_shot_iterator(dataset).get_next()
+ get_next = self.getNext(dataset)
for expected_features in self._next_expected_batch(
expected_output,
@@ -92,7 +90,7 @@
batch_size,
num_epochs,
):
- actual_features = self.evaluate(nxt)
+ actual_features = self.evaluate(get_next())
if label_name is not None:
expected_labels = expected_features.pop(label_name)
@@ -104,7 +102,7 @@
self.assertAllEqual(expected_features[k], actual_features[k])
with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(nxt)
+ self.evaluate(get_next())
def _test_dataset(self,
inputs,
@@ -118,18 +116,15 @@
# Convert str type because py3 tf strings are bytestrings
filenames = self._setup_files(
inputs, compression_type=kwargs.get("compression_type", None))
- with ops.Graph().as_default() as g:
- with self.session(graph=g) as sess:
- dataset = self._make_csv_dataset(
- filenames,
- batch_size=batch_size,
- num_epochs=num_epochs,
- label_name=label_name,
- **kwargs)
- self._verify_output(sess, dataset, batch_size, num_epochs, label_name,
- expected_output, expected_keys)
+ dataset = self._make_csv_dataset(
+ filenames,
+ batch_size=batch_size,
+ num_epochs=num_epochs,
+ label_name=label_name,
+ **kwargs)
+ self._verify_output(dataset, batch_size, num_epochs, label_name,
+ expected_output, expected_keys)
- @test_util.run_deprecated_v1
def testMakeCSVDataset(self):
"""Tests making a CSV dataset with keys and defaults provided."""
record_defaults = [
@@ -161,7 +156,6 @@
column_defaults=record_defaults,
)
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withBatchSizeAndEpochs(self):
"""Tests making a CSV dataset with keys and defaults provided."""
record_defaults = [
@@ -193,7 +187,6 @@
column_defaults=record_defaults,
)
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withCompressionType(self):
"""Tests `compression_type` argument."""
record_defaults = [
@@ -262,7 +255,6 @@
label_name="not_a_real_label",
column_names=column_names)
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withNoLabel(self):
"""Tests making a CSV dataset with no label provided."""
record_defaults = [
@@ -292,7 +284,6 @@
column_defaults=record_defaults,
)
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withNoHeader(self):
"""Tests that datasets can be created from CSV files with no header line.
"""
@@ -354,7 +345,6 @@
column_defaults=record_defaults,
)
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withNoColNames(self):
"""Tests that datasets can be created when column names are not specified.
@@ -459,7 +449,6 @@
header=True,
)
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withSelectCols(self):
record_defaults = [
constant_op.constant([], dtypes.int32),
@@ -566,7 +555,6 @@
label_name=None,
select_columns=["invalid_col_name"])
- @test_util.run_deprecated_v1
def testMakeCSVDataset_withShuffle(self):
record_defaults = [
constant_op.constant([], dtypes.int32),
@@ -591,69 +579,65 @@
total_records = 20
for batch_size in [1, 2]:
- with ops.Graph().as_default() as g:
- with self.session(graph=g) as sess:
- # Test that shuffling with the same seed produces the same result
- dataset1 = self._make_csv_dataset(
- filenames,
- column_defaults=record_defaults,
- column_names=column_names,
- batch_size=batch_size,
- header=True,
- shuffle=True,
- shuffle_seed=5,
- num_epochs=2,
- )
- dataset2 = self._make_csv_dataset(
- filenames,
- column_defaults=record_defaults,
- column_names=column_names,
- batch_size=batch_size,
- header=True,
- shuffle=True,
- shuffle_seed=5,
- num_epochs=2,
- )
- outputs1 = dataset_ops.make_one_shot_iterator(dataset1).get_next()
- outputs2 = dataset_ops.make_one_shot_iterator(dataset2).get_next()
- for _ in range(total_records // batch_size):
- batch1 = nest.flatten(self.evaluate(outputs1))
- batch2 = nest.flatten(self.evaluate(outputs2))
- for i in range(len(batch1)):
- self.assertAllEqual(batch1[i], batch2[i])
+ # Test that shuffling with the same seed produces the same result
+ dataset1 = self._make_csv_dataset(
+ filenames,
+ column_defaults=record_defaults,
+ column_names=column_names,
+ batch_size=batch_size,
+ header=True,
+ shuffle=True,
+ shuffle_seed=5,
+ num_epochs=2,
+ )
+ dataset2 = self._make_csv_dataset(
+ filenames,
+ column_defaults=record_defaults,
+ column_names=column_names,
+ batch_size=batch_size,
+ header=True,
+ shuffle=True,
+ shuffle_seed=5,
+ num_epochs=2,
+ )
+ next1 = self.getNext(dataset1)
+ next2 = self.getNext(dataset2)
+ for _ in range(total_records // batch_size):
+ batch1 = nest.flatten(self.evaluate(next1()))
+ batch2 = nest.flatten(self.evaluate(next2()))
+ for i in range(len(batch1)):
+ self.assertAllEqual(batch1[i], batch2[i])
- with ops.Graph().as_default() as g:
- with self.session(graph=g) as sess:
- # Test that shuffling with a different seed produces different results
- dataset1 = self._make_csv_dataset(
- filenames,
- column_defaults=record_defaults,
- column_names=column_names,
- batch_size=batch_size,
- header=True,
- shuffle=True,
- shuffle_seed=5,
- num_epochs=2,
- )
- dataset2 = self._make_csv_dataset(
- filenames,
- column_defaults=record_defaults,
- column_names=column_names,
- batch_size=batch_size,
- header=True,
- shuffle=True,
- shuffle_seed=6,
- num_epochs=2,
- )
- outputs1 = dataset_ops.make_one_shot_iterator(dataset1).get_next()
- outputs2 = dataset_ops.make_one_shot_iterator(dataset2).get_next()
- all_equal = False
- for _ in range(total_records // batch_size):
- batch1 = nest.flatten(self.evaluate(outputs1))
- batch2 = nest.flatten(self.evaluate(outputs2))
- for i in range(len(batch1)):
- all_equal = all_equal and np.array_equal(batch1[i], batch2[i])
- self.assertFalse(all_equal)
+ # Test that shuffling with a different seed produces different results
+ dataset1 = self._make_csv_dataset(
+ filenames,
+ column_defaults=record_defaults,
+ column_names=column_names,
+ batch_size=batch_size,
+ header=True,
+ shuffle=True,
+ shuffle_seed=5,
+ num_epochs=2,
+ )
+ dataset2 = self._make_csv_dataset(
+ filenames,
+ column_defaults=record_defaults,
+ column_names=column_names,
+ batch_size=batch_size,
+ header=True,
+ shuffle=True,
+ shuffle_seed=6,
+ num_epochs=2,
+ )
+ next1 = self.getNext(dataset1)
+ next2 = self.getNext(dataset2)
+ all_equal = False
+ for _ in range(total_records // batch_size):
+ batch1 = nest.flatten(self.evaluate(next1()))
+ batch2 = nest.flatten(self.evaluate(next2()))
+ for i in range(len(batch1)):
+ all_equal = all_equal and np.array_equal(batch1[i], batch2[i])
+ self.assertFalse(all_equal)
def testIndefiniteRepeatShapeInference(self):
column_names = ["col%d" % i for i in range(5)]
diff --git a/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py b/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py
index e6e24c3..a8a65dd 100644
--- a/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/map_and_batch_test.py
@@ -32,6 +32,7 @@
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 control_flow_util
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import script_ops
from tensorflow.python.platform import test
@@ -154,7 +155,7 @@
else:
self.assertEqual([None, 1], iterator.output_shapes.as_list())
next_element = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
self.assertAllEqual([[0], [1], [4], [9]], self.evaluate(next_element))
self.assertAllEqual([[16], [25], [36], [49]], self.evaluate(next_element))
if not drop_remainder:
@@ -180,7 +181,7 @@
iterator = dataset_ops.make_one_shot_iterator(dataset)
self.assertEqual([None, 1], iterator.output_shapes.as_list())
next_element = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
self.assertAllEqual([[0], [1], [4], [9]], self.evaluate(next_element))
self.assertAllEqual([[16], [25], [36], [49]], self.evaluate(next_element))
self.assertAllEqual([[64], [81]], self.evaluate(next_element))
@@ -204,7 +205,7 @@
elements = []
for _ in range(100):
elements.append(iterator.get_next())
- with self.cached_session() as sess:
+ with self.cached_session():
for i in range(5):
got = self.evaluate(elements)
got.sort(key=lambda x: x[0])
@@ -234,7 +235,7 @@
elements = []
for _ in range(100):
elements.append(iterator.get_next())
- with self.cached_session() as sess:
+ with self.cached_session():
for i in range(4):
got = self.evaluate(elements)
got.sort(key=lambda x: x[0])
@@ -267,7 +268,7 @@
init_op = iterator.initializer
get_next = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
self.evaluate(init_op)
for i in range(2):
actual = self.evaluate(get_next)
@@ -329,7 +330,7 @@
init_op = iterator.initializer
get_next = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
self.evaluate(init_op)
with self.assertRaisesRegexp(errors.InvalidArgumentError,
"number of elements does not match"):
@@ -361,7 +362,7 @@
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
for _ in range(3):
self.evaluate(get_next)
@@ -380,13 +381,11 @@
("6NUMA", 99, True),
)
@test_util.run_deprecated_v1
- def testMapAndBatchOutOfRangeError(self, threshold, numa_aware):
+ def testMapAndBatchMapError(self, threshold, numa_aware):
def raising_py_fn(i):
- if i == threshold:
+ if i >= threshold:
raise StopIteration()
- elif i > threshold:
- raise RuntimeError("Alternate error; you shouldn't see me! (i: %s)" % i)
else:
return i
@@ -401,14 +400,19 @@
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
for i in range(threshold // 10):
self.assertAllEqual([i * 10 + j for j in range(10)],
self.evaluate(get_next))
- if threshold % 10 != 0:
- self.assertAllEqual(
- [threshold // 10 * 10 + j for j in range(threshold % 10)],
- self.evaluate(get_next))
+ if numa_aware:
+ if threshold % 10 != 0:
+ self.assertAllEqual(
+ [threshold // 10 * 10 + j for j in range(threshold % 10)],
+ self.evaluate(get_next))
+ else:
+ for i in range(threshold // 10, 10):
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(get_next)
with self.assertRaises(errors.OutOfRangeError):
self.evaluate(get_next)
@@ -451,7 +455,7 @@
get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
for _ in range(10):
self.assertAllEqual([element for _ in range(10)],
self.evaluate(get_next))
@@ -497,10 +501,10 @@
def testMapAndBatchControlFlow(self, numa_aware):
def map_fn(x):
- previous_cond_v2_value = control_flow_ops.ENABLE_COND_V2
- control_flow_ops.ENABLE_COND_V2 = True
+ previous_control_flow_v2_value = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
return_value = control_flow_ops.cond(x < 50, lambda: x + 1, lambda: x * x)
- control_flow_ops.ENABLE_COND_V2 = previous_cond_v2_value
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = previous_control_flow_v2_value
return return_value
dataset = dataset_ops.Dataset.range(100).apply(
@@ -511,9 +515,8 @@
dataset = dataset.with_options(options)
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
- with self.cached_session() as sess:
+ with self.cached_session():
for i in range(10):
- print("Case %d" % i)
if i < 5:
self.assertAllEqual([i * 10 + j + 1 for j in range(10)],
self.evaluate(get_next))
diff --git a/tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py b/tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py
index 6042ca1..e41030d 100644
--- a/tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/map_defun_op_test.py
@@ -31,11 +31,11 @@
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import data_flow_ops
-from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
+# TODO(b/119837791): add eager coverage.
class MapDefunTest(test_base.DatasetTestBase):
def testMapDefunSimple(self):
@@ -254,46 +254,5 @@
self.assertAllEqual(self.evaluate(expected), self.evaluate(map_defun_op))
-class MapDefunBenchmark(test.Benchmark):
-
- def _run(self, op, name=None, num_iters=3000):
- with session.Session() as sess:
- # Warm up the session
- for _ in range(5):
- self.evaluate(op)
- start = time.time()
- for _ in range(num_iters):
- self.evaluate(op)
- end = time.time()
- mean_us = (end - start) * 1e6 / num_iters
- self.report_benchmark(
- name=name,
- iters=num_iters,
- wall_time=mean_us,
- extras={"examples_per_sec": num_iters / (end - start)})
-
- def benchmarkDefunVsMapFn(self):
- """Benchmarks to compare the performance of MapDefun vs tf.map_fn."""
-
- @function.defun(input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
- def defun(x):
- return array_ops.identity(x)
-
- def map_fn(x):
- return array_ops.identity(x)
-
- base = math_ops.range(100)
- for input_size in [10, 100, 1000, 10000]:
- num_iters = 100000 // input_size
- map_defun_op = map_defun.map_defun(defun, [base], [dtypes.int32], [()])
- map_fn_op = functional_ops.map_fn(map_fn, base)
-
- self._run(
- map_defun_op,
- "benchmarkMapDefun_size_%d" % input_size,
- num_iters=num_iters)
- self._run(
- map_fn_op, "benchmarkMapFn_size_%d" % input_size, num_iters=num_iters)
-
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py b/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py
index 0ee7616..fe83b4c 100644
--- a/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py
@@ -23,14 +23,14 @@
from tensorflow.python.data.experimental.ops import matching_files
from tensorflow.python.data.kernel_tests import test_base
-from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.util import compat
-class MatchingFilesTest(test_base.DatasetTestBase):
+@test_util.run_all_in_graph_and_eager_modes
+class MatchingFilesDatasetTest(test_base.DatasetTestBase):
def setUp(self):
self.tmp_dir = tempfile.mkdtemp()
@@ -42,30 +42,23 @@
for filename in filenames:
open(os.path.join(self.tmp_dir, filename), 'a').close()
- @test_util.run_deprecated_v1
def testNonExistingDirectory(self):
"""Test the MatchingFiles dataset with a non-existing directory."""
self.tmp_dir = os.path.join(self.tmp_dir, 'nonexistingdir')
dataset = matching_files.MatchingFilesDataset(
os.path.join(self.tmp_dir, '*'))
- with self.cached_session() as sess:
- next_element = dataset_ops.make_one_shot_iterator(dataset).get_next()
- with self.assertRaises(errors.NotFoundError):
- sess.run(next_element)
+ self.assertDatasetProduces(
+ dataset, expected_error=(errors.NotFoundError, ''))
- @test_util.run_deprecated_v1
def testEmptyDirectory(self):
"""Test the MatchingFiles dataset with an empty directory."""
dataset = matching_files.MatchingFilesDataset(
os.path.join(self.tmp_dir, '*'))
- with self.cached_session() as sess:
- next_element = dataset_ops.make_one_shot_iterator(dataset).get_next()
- with self.assertRaises(errors.NotFoundError):
- sess.run(next_element)
+ self.assertDatasetProduces(
+ dataset, expected_error=(errors.NotFoundError, ''))
- @test_util.run_deprecated_v1
def testSimpleDirectory(self):
"""Test the MatchingFiles dataset with a simple directory."""
@@ -74,21 +67,14 @@
dataset = matching_files.MatchingFilesDataset(
os.path.join(self.tmp_dir, '*'))
- with self.cached_session() as sess:
- next_element = dataset_ops.make_one_shot_iterator(dataset).get_next()
+ self.assertDatasetProduces(
+ dataset,
+ expected_output=[
+ compat.as_bytes(os.path.join(self.tmp_dir, filename))
+ for filename in filenames
+ ],
+ assert_items_equal=True)
- expected_filenames = []
- actual_filenames = []
- for filename in filenames:
- expected_filenames.append(
- compat.as_bytes(os.path.join(self.tmp_dir, filename)))
- actual_filenames.append(compat.as_bytes(sess.run(next_element)))
-
- self.assertItemsEqual(expected_filenames, actual_filenames)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_element)
-
- @test_util.run_deprecated_v1
def testFileSuffixes(self):
"""Test the MatchingFiles dataset using the suffixes of filename."""
@@ -97,20 +83,14 @@
dataset = matching_files.MatchingFilesDataset(
os.path.join(self.tmp_dir, '*.py'))
- with self.cached_session() as sess:
- next_element = dataset_ops.make_one_shot_iterator(dataset).get_next()
- expected_filenames = []
- actual_filenames = []
- for filename in filenames[1:-1]:
- expected_filenames.append(
- compat.as_bytes(os.path.join(self.tmp_dir, filename)))
- actual_filenames.append(compat.as_bytes(sess.run(next_element)))
+ self.assertDatasetProduces(
+ dataset,
+ expected_output=[
+ compat.as_bytes(os.path.join(self.tmp_dir, filename))
+ for filename in filenames[1:-1]
+ ],
+ assert_items_equal=True)
- self.assertItemsEqual(expected_filenames, actual_filenames)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_element)
-
- @test_util.run_deprecated_v1
def testFileMiddles(self):
"""Test the MatchingFiles dataset using the middles of filename."""
@@ -119,20 +99,14 @@
dataset = matching_files.MatchingFilesDataset(
os.path.join(self.tmp_dir, 'b*.py*'))
- with self.cached_session() as sess:
- next_element = dataset_ops.make_one_shot_iterator(dataset).get_next()
- expected_filenames = []
- actual_filenames = []
- for filename in filenames[1:3]:
- expected_filenames.append(
- compat.as_bytes(os.path.join(self.tmp_dir, filename)))
- actual_filenames.append(compat.as_bytes(sess.run(next_element)))
+ self.assertDatasetProduces(
+ dataset,
+ expected_output=[
+ compat.as_bytes(os.path.join(self.tmp_dir, filename))
+ for filename in filenames[1:3]
+ ],
+ assert_items_equal=True)
- self.assertItemsEqual(expected_filenames, actual_filenames)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_element)
-
- @test_util.run_deprecated_v1
def testNestedDirectories(self):
"""Test the MatchingFiles dataset with nested directories."""
@@ -156,21 +130,20 @@
]
dataset = matching_files.MatchingFilesDataset(patterns)
- with self.cached_session() as sess:
- next_element = dataset_ops.make_one_shot_iterator(dataset).get_next()
- expected_filenames = [
- compat.as_bytes(filename)
- for filename in filenames
- if filename.endswith('.txt') or filename.endswith('.log')
- ]
- actual_filenames = []
- while True:
- try:
- actual_filenames.append(compat.as_bytes(sess.run(next_element)))
- except errors.OutOfRangeError:
- break
+ next_element = self.getNext(dataset)
+ expected_filenames = [
+ compat.as_bytes(filename)
+ for filename in filenames
+ if filename.endswith('.txt') or filename.endswith('.log')
+ ]
+ actual_filenames = []
+ while True:
+ try:
+ actual_filenames.append(compat.as_bytes(self.evaluate(next_element())))
+ except errors.OutOfRangeError:
+ break
- self.assertItemsEqual(expected_filenames, actual_filenames)
+ self.assertItemsEqual(expected_filenames, actual_filenames)
if __name__ == '__main__':
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD
index f214944..bf868eb 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/BUILD
@@ -279,7 +279,6 @@
"//tensorflow/python:errors",
"//tensorflow/python:math_ops",
"//tensorflow/python/data/experimental/ops:optimization",
- "//tensorflow/python/data/experimental/ops:optimization_options",
"//tensorflow/python/data/kernel_tests:test_base",
"//tensorflow/python/data/ops:dataset_ops",
],
@@ -287,7 +286,7 @@
py_test(
name = "optimize_dataset_test",
- size = "small",
+ size = "medium",
srcs = ["optimize_dataset_test.py"],
srcs_version = "PY2AND3",
tags = [
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py
index 7371cf3..3ce921b 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py
@@ -20,7 +20,6 @@
from absl.testing import parameterized
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
@@ -72,7 +71,6 @@
dataset = dataset.cache()
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.filter_fusion = True
dataset = dataset.with_options(options)
expected_output = []
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py
index 5f3a868..f080891 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/hoist_random_uniform_test.py
@@ -20,7 +20,6 @@
from absl.testing import parameterized
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import context
@@ -92,7 +91,6 @@
["Zip[0]", "Map"] if will_optimize else ["Map"])).map(function)
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.hoist_random_uniform = True
dataset = dataset.with_options(options)
self._testDataset(dataset)
@@ -109,7 +107,6 @@
dataset = dataset_ops.Dataset.range(5).apply(
optimization.assert_next(["Zip[0]", "Map"])).map(random_with_capture)
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.hoist_random_uniform = True
dataset = dataset.with_options(options)
self._testDataset(dataset)
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py
index fc65f52..8af86da 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/latency_all_edges_test.py
@@ -20,7 +20,6 @@
from tensorflow.python.data.experimental.kernel_tests import stats_dataset_test_base
from tensorflow.python.data.experimental.ops import optimization
from tensorflow.python.data.experimental.ops import stats_aggregator
-from tensorflow.python.data.experimental.ops import stats_options
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
@@ -36,7 +35,6 @@
["LatencyStats", "Map", "LatencyStats", "Prefetch",
"LatencyStats"])).map(lambda x: x * x).prefetch(1)
options = dataset_ops.Options()
- options.experimental_stats = stats_options.StatsOptions()
options.experimental_stats.latency_all_edges = True
options.experimental_stats.aggregator = aggregator
dataset = dataset.with_options(options)
@@ -53,29 +51,6 @@
self._assertSummaryHasCount(summary_str,
"record_latency_PrefetchDataset/_6", 1)
- def testLatencyStatsOptimizationV2(self):
- aggregator = stats_aggregator.StatsAggregator()
- dataset = dataset_ops.Dataset.from_tensors(1).apply(
- optimization.assert_next(
- ["LatencyStats", "Map", "LatencyStats", "Prefetch",
- "LatencyStats"])).map(lambda x: x * x).prefetch(1)
- options = dataset_ops.Options()
- options.experimental_stats = stats_options.StatsOptions()
- options.experimental_stats.aggregator = aggregator
- dataset = dataset.with_options(options)
- self.assertDatasetProduces(
- dataset,
- expected_output=[1],
- requires_initialization=True,
- num_test_iterations=1)
- summary_t = aggregator.get_summary()
- summary_str = self.evaluate(summary_t)
- self._assertSummaryHasCount(summary_str, "record_latency_TensorDataset/_1",
- 1)
- self._assertSummaryHasCount(summary_str, "record_latency_MapDataset/_4", 1)
- self._assertSummaryHasCount(summary_str,
- "record_latency_PrefetchDataset/_6", 1)
-
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py
index 801f664..e2ff311 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py
@@ -18,7 +18,6 @@
from __future__ import print_function
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import test_util
@@ -32,10 +31,6 @@
dataset = dataset_ops.Dataset.range(10).apply(
optimization.assert_next(
["MapAndBatch"])).map(lambda x: x * x).batch(10)
- options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
- options.experimental_optimization.map_and_batch_fusion = True
- dataset = dataset.with_options(options)
self.assertDatasetProduces(
dataset, expected_output=[[x * x for x in range(10)]])
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py
index db8f214..fa1d673 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py
@@ -20,7 +20,6 @@
from absl.testing import parameterized
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
@@ -84,7 +83,6 @@
optimization.assert_next(
["Map", "FilterByLastComponent"])).map(function).filter(predicate)
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.map_and_filter_fusion = True
dataset = dataset.with_options(options)
self._testMapAndFilter(dataset, function, predicate)
@@ -103,7 +101,6 @@
optimization.assert_next(["Map",
"Filter"])).map(function).filter(predicate)
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.map_and_filter_fusion = True
dataset = dataset.with_options(options)
self._testMapAndFilter(dataset, function, predicate)
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py
index d8d6390..defdaf0 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py
@@ -20,7 +20,6 @@
from absl.testing import parameterized
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import test_util
@@ -75,7 +74,6 @@
dataset = dataset.cache()
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.map_fusion = True
dataset = dataset.with_options(options)
expected_output = []
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py
index 0ff3fff..d8dd31f 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py
@@ -20,7 +20,6 @@
from absl.testing import parameterized
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
@@ -68,7 +67,6 @@
dataset = dataset_ops.Dataset.range(5).apply(
optimization.assert_next(next_nodes)).map(function)
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.map_parallelization = True
dataset = dataset.with_options(options)
if should_optimize:
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py
index c2e08e2..65fa2ba 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/map_vectorization_test.py
@@ -23,7 +23,6 @@
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
@@ -344,17 +343,20 @@
Tuple of (unoptimized dataset, optimized dataset).
"""
map_node_name = "Map" if num_parallel_calls is None else "ParallelMap"
- batch_size = 100
def _make_dataset(node_names):
- return base_dataset.apply(optimization.assert_next(node_names)).map(
- map_fn, num_parallel_calls=num_parallel_calls).batch(batch_size)
+ dataset = base_dataset.apply(optimization.assert_next(node_names))
+ dataset = dataset.map(map_fn, num_parallel_calls)
+ dataset = dataset.batch(100)
+ options = dataset_ops.Options()
+ options.experimental_optimization.map_and_batch_fusion = False
+ dataset = dataset.with_options(options)
+ return dataset
unoptimized = _make_dataset([map_node_name, "Batch"])
optimized = _make_dataset(["Batch", map_node_name]
if expect_optimized else [map_node_name, "Batch"])
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.map_vectorization = True
optimized = optimized.with_options(options)
return unoptimized, optimized
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py
index ce86bfa..8058f53 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py
@@ -18,7 +18,6 @@
from __future__ import print_function
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
@@ -42,10 +41,6 @@
["FiniteRepeat", "FiniteSkip", "Prefetch", "MemoryCacheImpl"]))
dataset = dataset.repeat(some_tensor).skip(5).take(-1).skip(0).repeat(
1).prefetch(0).prefetch(1).cache()
- options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
- options.experimental_optimization.noop_elimination = True
- dataset = dataset.with_options(options)
self.assertDatasetProduces(dataset, expected_output=range(5))
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_test.py
index 751be83..dd432b8 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/optimize_dataset_test.py
@@ -25,7 +25,6 @@
from tensorflow.python.data.experimental.ops import batching
from tensorflow.python.data.experimental.ops import grouping
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops import optimization_options
from tensorflow.python.data.experimental.ops import scan_ops
from tensorflow.python.data.experimental.ops import threadpool
from tensorflow.python.data.kernel_tests import test_base
@@ -163,7 +162,13 @@
dataset = dataset_ops.Dataset.range(1)
dataset = dataset.flat_map(flat_map_fn)
- dataset = dataset_ops._OptimizeDataset(dataset, ["map_and_batch_fusion"])
+
+ # TODO(b/120558523): We use Options instead of _OptimizeDataset directly
+ # here because of a bug with chaining _OptimizeDatasets when there are
+ # nested dataset functions
+ options = dataset_ops.Options()
+ options.experimental_optimization.map_and_batch_fusion = True
+ dataset = dataset.with_options(options)
self.assertDatasetProduces(dataset, expected_output=[[0]])
def testOptimizationThreadPoolDataset(self):
@@ -209,10 +214,8 @@
unoptimized_dataset = dataset_fn(variable)
options = dataset_ops.Options()
- opt_options = optimization_options.OptimizationOptions()
- opt_options.noop_elimination = True
- opt_options.map_and_batch_fusion = True
- options.experimental_optimization = opt_options
+ options.experimental_optimization.noop_elimination = True
+ options.experimental_optimization.map_and_batch_fusion = True
optimized_dataset = unoptimized_dataset.with_options(options)
# Check that warning is logged.
@@ -222,10 +225,10 @@
self.assertGreaterEqual(len(w), 1)
expected = ("tf.data static optimizations are not compatible with "
- "tf.Variable. The following optimizations will be disabled: "
- "map_and_batch_fusion, noop_elimination. To enable "
- "optimizations, use resource variables instead by calling "
- "`tf.enable_resource_variables()` at the start of the program.")
+ "tf.Variable. The following optimizations will be disabled: %s."
+ " To enable optimizations, use resource variables instead by "
+ "calling `tf.enable_resource_variables()` at the start of the "
+ "program." % (", ".join(options._static_optimizations())))
self.assertTrue(any([expected in str(warning) for warning in w]))
# Check that outputs are the same in the optimized and unoptimized cases,
@@ -245,6 +248,29 @@
except errors.OutOfRangeError:
break
+ def testOptimizationEnabledByDefault(self):
+ """Tests that some optimizations are applied to datasets by default."""
+ options = dataset_ops.Options()
+ expected_optimizations = [
+ "map_and_batch_fusion",
+ "noop_elimination",
+ "shuffle_and_repeat_fusion",
+ ]
+ self.assertEqual(
+ set(options._static_optimizations()), set(expected_optimizations))
+
+ def testOptimizationDisableDefault(self):
+ """Tests that we can disable all static optimizations enabled by default.
+
+ If the `apply_default_optimizations` optimization options flag is False,
+ only explicitly enabled optimizations will be applied.
+ """
+ options = dataset_ops.Options()
+ options.experimental_optimization.hoist_random_uniform = True
+ options.experimental_optimization.apply_default_optimizations = False
+ expected_optimizations = ["hoist_random_uniform"]
+ self.assertEqual(options._static_optimizations(), expected_optimizations)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py b/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py
index 5f746ec..594b593 100644
--- a/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py
@@ -18,7 +18,6 @@
from __future__ import print_function
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
@@ -32,10 +31,6 @@
def testShuffleAndRepeatFusion(self):
dataset = dataset_ops.Dataset.range(10).apply(
optimization.assert_next(["ShuffleAndRepeat"])).shuffle(10).repeat(2)
- options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
- options.experimental_optimization.shuffle_and_repeat_fusion = True
- dataset = dataset.with_options(options)
get_next = self.getNext(dataset)
for _ in range(2):
diff --git a/tensorflow/python/data/experimental/kernel_tests/override_threadpool_test.py b/tensorflow/python/data/experimental/kernel_tests/override_threadpool_test.py
index aa81663..811a582 100644
--- a/tensorflow/python/data/experimental/kernel_tests/override_threadpool_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/override_threadpool_test.py
@@ -35,6 +35,7 @@
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class OverrideThreadpoolTest(test_base.DatasetTestBase,
parameterized.TestCase):
@@ -53,14 +54,12 @@
lambda x: script_ops.py_func(get_thread_id, [x], dtypes.int64),
num_parallel_calls=32).apply(unique.unique()))
dataset = override_threadpool_fn(dataset)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
+ next_element = self.getNext(dataset, requires_initialization=True)
- self.evaluate(iterator.initializer)
thread_ids = []
try:
while True:
- thread_ids.append(self.evaluate(next_element))
+ thread_ids.append(self.evaluate(next_element()))
except errors.OutOfRangeError:
pass
self.assertLen(thread_ids, len(set(thread_ids)))
@@ -82,7 +81,6 @@
("8", 4, 1),
("9", 4, 4),
)
- @test_util.run_deprecated_v1
def testNumThreadsDeprecated(self, num_threads, max_intra_op_parallelism):
def override_threadpool_fn(dataset):
@@ -109,7 +107,6 @@
("11", 4, 4),
("12", None, None),
)
- @test_util.run_deprecated_v1
def testNumThreads(self, num_threads, max_intra_op_parallelism):
def override_threadpool_fn(dataset):
diff --git a/tensorflow/python/data/experimental/kernel_tests/prefetch_to_device_test.py b/tensorflow/python/data/experimental/kernel_tests/prefetch_to_device_test.py
index 80bd43e..2af31ad 100644
--- a/tensorflow/python/data/experimental/kernel_tests/prefetch_to_device_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/prefetch_to_device_test.py
@@ -29,6 +29,7 @@
from tensorflow.python.platform import test
+# TODO(b/119837791): add eager coverage when supported.
class PrefetchToDeviceTest(test_base.DatasetTestBase):
@test_util.run_deprecated_v1
diff --git a/tensorflow/python/data/experimental/kernel_tests/rejection_resample_test.py b/tensorflow/python/data/experimental/kernel_tests/rejection_resample_test.py
index 76f68f5..4d35b16 100644
--- a/tensorflow/python/data/experimental/kernel_tests/rejection_resample_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/rejection_resample_test.py
@@ -17,11 +17,9 @@
from __future__ import division
from __future__ import print_function
-import time
from absl.testing import parameterized
import numpy as np
-from six.moves import xrange # pylint: disable=redefined-builtin
from tensorflow.python.data.experimental.ops import resampling
from tensorflow.python.data.kernel_tests import test_base
@@ -36,35 +34,12 @@
from tensorflow.python.util import compat
-def _time_resampling(
- test_obj, data_np, target_dist, init_dist, num_to_sample):
- dataset = dataset_ops.Dataset.from_tensor_slices(data_np).repeat()
-
- # Reshape distribution via rejection sampling.
- dataset = dataset.apply(
- resampling.rejection_resample(
- class_func=lambda x: x,
- target_dist=target_dist,
- initial_dist=init_dist,
- seed=142))
-
- get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
-
- with test_obj.test_session() as sess:
- start_time = time.time()
- for _ in xrange(num_to_sample):
- sess.run(get_next)
- end_time = time.time()
-
- return end_time - start_time
-
-
+@test_util.run_all_in_graph_and_eager_modes
class RejectionResampleTest(test_base.DatasetTestBase, parameterized.TestCase):
@parameterized.named_parameters(
("InitialDistributionKnown", True),
("InitialDistributionUnknown", False))
- @test_util.run_deprecated_v1
def testDistribution(self, initial_known):
classes = np.random.randint(5, size=(20000,)) # Uniformly sampled
target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
@@ -73,17 +48,17 @@
dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
200, seed=21).map(lambda c: (c, string_ops.as_string(c))).repeat()
- get_next = dataset_ops.make_one_shot_iterator(dataset.apply(
- resampling.rejection_resample(
- target_dist=target_dist,
- initial_dist=initial_dist,
- class_func=lambda c, _: c,
- seed=27))).get_next()
+ get_next = self.getNext(
+ dataset.apply(
+ resampling.rejection_resample(
+ target_dist=target_dist,
+ initial_dist=initial_dist,
+ class_func=lambda c, _: c,
+ seed=27)))
- with self.cached_session() as sess:
- returned = []
- while len(returned) < 4000:
- returned.append(sess.run(get_next))
+ returned = []
+ while len(returned) < 4000:
+ returned.append(self.evaluate(get_next()))
returned_classes, returned_classes_and_data = zip(*returned)
_, returned_data = zip(*returned_classes_and_data)
@@ -99,7 +74,6 @@
@parameterized.named_parameters(
("OnlyInitial", True),
("NotInitial", False))
- @test_util.run_deprecated_v1
def testEdgeCasesSampleFromInitialDataset(self, only_initial_dist):
init_dist = [0.5, 0.5]
target_dist = [0.5, 0.5] if only_initial_dist else [0.0, 1.0]
@@ -117,15 +91,13 @@
target_dist=target_dist,
initial_dist=init_dist))
- get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- returned = []
- with self.assertRaises(errors.OutOfRangeError):
- while True:
- returned.append(sess.run(get_next))
+ returned = []
+ with self.assertRaises(errors.OutOfRangeError):
+ while True:
+ returned.append(self.evaluate(get_next()))
- @test_util.run_deprecated_v1
def testRandomClasses(self):
init_dist = [0.25, 0.25, 0.25, 0.25]
target_dist = [0.0, 0.0, 0.0, 1.0]
@@ -149,13 +121,12 @@
target_dist=target_dist,
initial_dist=init_dist))
- get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- returned = []
- with self.assertRaises(errors.OutOfRangeError):
- while True:
- returned.append(sess.run(get_next))
+ returned = []
+ with self.assertRaises(errors.OutOfRangeError):
+ while True:
+ returned.append(self.evaluate(get_next()))
classes, _ = zip(*returned)
bincount = np.bincount(
@@ -165,22 +136,5 @@
self.assertAllClose(target_dist, bincount, atol=1e-2)
-class ResampleDatasetBenchmark(test.Benchmark):
-
- def benchmarkResamplePerformance(self):
- init_dist = [0.25, 0.25, 0.25, 0.25]
- target_dist = [0.0, 0.0, 0.0, 1.0]
- num_classes = len(init_dist)
- # We don't need many samples to test a dirac-delta target distribution
- num_samples = 1000
- data_np = np.random.choice(num_classes, num_samples, p=init_dist)
-
- resample_time = _time_resampling(
- self, data_np, target_dist, init_dist, num_to_sample=1000)
-
- self.report_benchmark(
- iters=1000, wall_time=resample_time, name="benchmark_resample")
-
-
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/restructured_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/restructured_dataset_test.py
index 658e612..3b0d23d 100644
--- a/tensorflow/python/data/experimental/kernel_tests/restructured_dataset_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/restructured_dataset_test.py
@@ -27,6 +27,7 @@
from tensorflow.python.platform import test
+# TODO(b/119837791): Add eager coverage
class RestructuredDatasetTest(test_base.DatasetTestBase):
@test_util.run_deprecated_v1
diff --git a/tensorflow/python/data/experimental/kernel_tests/scan_test.py b/tensorflow/python/data/experimental/kernel_tests/scan_test.py
index 03af7ec..89b3824 100644
--- a/tensorflow/python/data/experimental/kernel_tests/scan_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/scan_test.py
@@ -24,58 +24,44 @@
from tensorflow.python.data.experimental.ops import scan_ops
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import script_ops
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class ScanTest(test_base.DatasetTestBase):
def _counting_dataset(self, start, scan_fn):
return dataset_ops.Dataset.from_tensors(0).repeat().apply(
scan_ops.scan(start, scan_fn))
- @test_util.run_deprecated_v1
def testCount(self):
def make_scan_fn(step):
return lambda state, _: (state + step, state)
- start = array_ops.placeholder(dtypes.int32, shape=[])
- step = array_ops.placeholder(dtypes.int32, shape=[])
- take = array_ops.placeholder(dtypes.int64, shape=[])
- iterator = dataset_ops.make_initializable_iterator(self._counting_dataset(
- start, make_scan_fn(step)).take(take))
- next_element = iterator.get_next()
+ def dataset_fn(start, step, take):
+ return self._counting_dataset(start, make_scan_fn(step)).take(take)
- with self.cached_session() as sess:
+ for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
+ (10, 2, 10), (10, -1, 10), (10, -2,
+ 10)]:
+ next_element = self.getNext(dataset_fn(start_val, step_val, take_val))
+ for expected, _ in zip(
+ itertools.count(start_val, step_val), range(take_val)):
+ self.assertEqual(expected, self.evaluate(next_element()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
- for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
- (10, 2, 10), (10, -1, 10),
- (10, -2, 10)]:
- sess.run(iterator.initializer,
- feed_dict={start: start_val, step: step_val, take: take_val})
- for expected, _ in zip(
- itertools.count(start_val, step_val), range(take_val)):
- self.assertEqual(expected, self.evaluate(next_element))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
-
- @test_util.run_in_graph_and_eager_modes
def testFibonacci(self):
- iterator = dataset_ops.make_one_shot_iterator(
- dataset_ops.Dataset.from_tensors(1).repeat(None).apply(
- scan_ops.scan([0, 1], lambda a, _: ([a[1], a[0] + a[1]], a[1]))))
-
- if context.executing_eagerly():
- next_element = iterator.get_next
- else:
- get_next = iterator.get_next()
- next_element = lambda: get_next
+ data = dataset_ops.Dataset.from_tensors(1).repeat(None).apply(
+ scan_ops.scan([0, 1], lambda a, _: ([a[1], a[0] + a[1]], a[1])))
+ next_element = self.getNext(data)
self.assertEqual(1, self.evaluate(next_element()))
self.assertEqual(1, self.evaluate(next_element()))
@@ -84,8 +70,10 @@
self.assertEqual(5, self.evaluate(next_element()))
self.assertEqual(8, self.evaluate(next_element()))
+ # TODO(b/119837791): Add coverage for eager.
@test_util.run_deprecated_v1
- def testSparseCount(self):
+ def testSkipEagerSparseCount(self):
+
def _sparse(i):
return sparse_tensor.SparseTensorValue(
indices=np.array([[0, 0]]),
@@ -95,27 +83,20 @@
def make_scan_fn(step):
return lambda state, _: (_sparse(state.values[0] + step), state)
- start = array_ops.placeholder(dtypes.int32, shape=[])
- step = array_ops.placeholder(dtypes.int32, shape=[])
- take = array_ops.placeholder(dtypes.int64, shape=[])
- iterator = dataset_ops.make_initializable_iterator(self._counting_dataset(
- _sparse(start), make_scan_fn(step)).take(take))
- next_element = iterator.get_next()
+ def dataset_fn(start, step, take):
+ return self._counting_dataset(_sparse(start),
+ make_scan_fn(step)).take(take)
- with self.cached_session() as sess:
+ for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
+ (10, 2, 10), (10, -1, 10), (10, -2,
+ 10)]:
+ next_element = self.getNext(dataset_fn(start_val, step_val, take_val))
+ for expected, _ in zip(
+ itertools.count(start_val, step_val), range(take_val)):
+ self.assertEqual(expected, self.evaluate(next_element()).values[0])
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
- for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
- (10, 2, 10), (10, -1, 10),
- (10, -2, 10)]:
- sess.run(iterator.initializer,
- feed_dict={start: start_val, step: step_val, take: take_val})
- for expected, _ in zip(
- itertools.count(start_val, step_val), range(take_val)):
- self.assertEqual(expected, self.evaluate(next_element).values[0])
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
-
- @test_util.run_deprecated_v1
def testChangingStateShape(self):
# Test the fixed-point shape invariant calculations: start with
# initial values with known shapes, and use a scan function that
@@ -133,16 +114,14 @@
self.assertIs(None, dataset.output_shapes[0][1].ndims)
self.assertEqual([], dataset.output_shapes[1].as_list())
- iterator = dataset_ops.make_one_shot_iterator(dataset)
- next_element = iterator.get_next()
+ next_element = self.getNext(dataset)
- with self.cached_session() as sess:
- for i in range(5):
- (longer_vector_val, larger_rank_val), _ = self.evaluate(next_element)
- self.assertAllEqual([0] * (2**i), longer_vector_val)
- self.assertAllEqual(np.array(1, ndmin=i), larger_rank_val)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(5):
+ (longer_vector_val, larger_rank_val), _ = self.evaluate(next_element())
+ self.assertAllEqual([0] * (2**i), longer_vector_val)
+ self.assertAllEqual(np.array(1, ndmin=i), larger_rank_val)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
def testIncorrectStateType(self):
@@ -169,6 +148,21 @@
dataset.apply(
scan_ops.scan(constant_op.constant(1, dtype=dtypes.int32), _scan_fn))
+ def testPreserveCardinality(self):
+
+ def scan_fn(state, val):
+
+ def py_fn(_):
+ raise StopIteration()
+
+ return state, script_ops.py_func(py_fn, [val], dtypes.int64)
+
+ dataset = dataset_ops.Dataset.from_tensors(0).apply(
+ scan_ops.scan(constant_op.constant(1), scan_fn))
+ get_next = self.getNext(dataset)
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(get_next())
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py
index e65aa44..bdbd870 100644
--- a/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py
@@ -23,6 +23,7 @@
import numpy as np
from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import dtypes
@@ -74,23 +75,39 @@
Raises:
AssertionError if any test fails.
"""
+ # NOTE: We disable all default optimizations in serialization tests in order
+ # to test the actual dataset in question.
+ options = dataset_ops.Options()
+ options.experimental_optimization = OptimizationOptions()
+ options.experimental_optimization.apply_default_optimizations = False
+
+ def ds_fn1_no_opt():
+ return ds_fn1().with_options(options)
+
self.verify_unused_iterator(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
self.verify_fully_used_iterator(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
self.verify_exhausted_iterator(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
self.verify_init_before_restore(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
self.verify_multiple_breaks(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
self.verify_reset_restored_iterator(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
self.verify_restore_in_empty_graph(
- ds_fn1, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt, num_outputs, sparse_tensors=sparse_tensors)
if ds_fn2:
+
+ def ds_fn2_no_opt():
+ return ds_fn2().with_options(options)
+
self.verify_restore_in_modified_graph(
- ds_fn1, ds_fn2, num_outputs, sparse_tensors=sparse_tensors)
+ ds_fn1_no_opt,
+ ds_fn2_no_opt,
+ num_outputs,
+ sparse_tensors=sparse_tensors)
def verify_unused_iterator(self,
ds_fn,
diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/map_and_batch_dataset_serialization_test.py b/tensorflow/python/data/experimental/kernel_tests/serialization/map_and_batch_dataset_serialization_test.py
index 166ffa9..8bfe6ce 100644
--- a/tensorflow/python/data/experimental/kernel_tests/serialization/map_and_batch_dataset_serialization_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/serialization/map_and_batch_dataset_serialization_test.py
@@ -22,6 +22,7 @@
from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base
from tensorflow.python.data.experimental.ops import batching
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
@@ -83,6 +84,19 @@
self.run_core_tests(lambda: build_ds(10, True), lambda: build_ds(15, True),
num_outputs_drop_remainder)
+ def testSparse(self):
+
+ def build_dataset():
+
+ def map_fn(i):
+ return sparse_tensor.SparseTensorValue(
+ indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+ return dataset_ops.Dataset.range(10).apply(
+ batching.map_and_batch(map_fn, 5))
+
+ self.run_core_tests(build_dataset, None, 2)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/experimental/kernel_tests/shuffle_and_repeat_test.py b/tensorflow/python/data/experimental/kernel_tests/shuffle_and_repeat_test.py
index 9528f83..110966a 100644
--- a/tensorflow/python/data/experimental/kernel_tests/shuffle_and_repeat_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/shuffle_and_repeat_test.py
@@ -23,11 +23,11 @@
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
-from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
class ShuffleAndRepeatTest(test_base.DatasetTestBase):
def _build_ds(self, seed, count=5, num_elements=20):
@@ -35,17 +35,15 @@
shuffle_ops.shuffle_and_repeat(buffer_size=5, count=count, seed=seed))
def _gen_outputs(self, ds_fn, num_outputs, verify_exhausted=True):
- get_next = dataset_ops.make_one_shot_iterator(ds_fn()).get_next()
+ get_next = self.getNext(ds_fn())
outputs = []
- with self.cached_session() as sess:
- for _ in range(num_outputs):
- outputs.append(self.evaluate(get_next))
- if verify_exhausted:
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(get_next)
+ for _ in range(num_outputs):
+ outputs.append(self.evaluate(get_next()))
+ if verify_exhausted:
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
return outputs
- @test_util.run_deprecated_v1
def testCorrectOutput(self):
output = self._gen_outputs(lambda: self._build_ds(10), 100)
self.assertSequenceEqual(
@@ -54,7 +52,6 @@
for i in range(5):
self.assertSequenceEqual(sorted(output[i * 20:(i + 1) * 20]), range(20))
- @test_util.run_deprecated_v1
def testReshuffling(self):
# Check that the output orders of different epochs are indeed different.
output = self._gen_outputs(lambda: self._build_ds(10), 100)
@@ -63,20 +60,17 @@
epoch2 = output[(i + 1) * 20:(i + 2) * 20]
self.assertNotEqual(epoch1, epoch2)
- @test_util.run_deprecated_v1
def testSameOrderForSameSeeds(self):
output1 = self._gen_outputs(lambda: self._build_ds(10), 100)
output2 = self._gen_outputs(lambda: self._build_ds(10), 100)
self.assertEqual(output1, output2)
- @test_util.run_deprecated_v1
def testDifferentOrderForDifferentSeeds(self):
output1 = self._gen_outputs(lambda: self._build_ds(10), 100)
output2 = self._gen_outputs(lambda: self._build_ds(20), 100)
self.assertNotEqual(output1, output2)
self.assertEqual(sorted(output1), sorted(output2))
- @test_util.run_deprecated_v1
def testCountNone(self):
output1 = self._gen_outputs(
lambda: self._build_ds(10, count=None), 100, verify_exhausted=False)
@@ -85,7 +79,6 @@
self.assertNotEqual(output1, output2)
self.assertEqual(sorted(output1), sorted(output2))
- @test_util.run_deprecated_v1
def testCountMinusOne(self):
output1 = self._gen_outputs(
lambda: self._build_ds(10, count=-1), 100, verify_exhausted=False)
@@ -110,12 +103,10 @@
100)
def testLargeBufferSize(self):
- with ops.Graph().as_default() as g:
- ds = dataset_ops.Dataset.range(20).apply(
- shuffle_ops.shuffle_and_repeat(buffer_size=21))
- get_next_op = ds.make_one_shot_iterator().get_next()
- with self.session(graph=g) as sess:
- self.evaluate(get_next_op)
+ ds = dataset_ops.Dataset.range(20).apply(
+ shuffle_ops.shuffle_and_repeat(buffer_size=21))
+ get_next = self.getNext(ds)
+ self.evaluate(get_next())
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/sleep_test.py b/tensorflow/python/data/experimental/kernel_tests/sleep_test.py
index 46b22f8..a4fe847 100644
--- a/tensorflow/python/data/experimental/kernel_tests/sleep_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/sleep_test.py
@@ -29,25 +29,21 @@
_NUMPY_RANDOM_SEED = 42
+@test_util.run_all_in_graph_and_eager_modes
class SleepTest(test_base.DatasetTestBase):
- @test_util.run_deprecated_v1
def testSleep(self):
sleep_microseconds = 100
dataset = dataset_ops.Dataset.range(10).apply(
sleep.sleep(sleep_microseconds))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
-
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- start_time = time.time()
- for i in range(10):
- self.assertEqual(i, self.evaluate(next_element))
- end_time = time.time()
- self.assertGreater(end_time - start_time, (10 * sleep_microseconds) / 1e6)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ next_element = self.getNext(dataset)
+ start_time = time.time()
+ for i in range(10):
+ self.assertEqual(i, self.evaluate(next_element()))
+ end_time = time.time()
+ self.assertGreater(end_time - start_time, (10 * sleep_microseconds) / 1e6)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py
index 8a30036..59d0ebd 100644
--- a/tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_ops_test.py
@@ -26,7 +26,6 @@
from tensorflow.python.data.experimental.ops import optimization
from tensorflow.python.data.experimental.ops import stats_aggregator
from tensorflow.python.data.experimental.ops import stats_ops
-from tensorflow.python.data.experimental.ops import stats_options
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
@@ -36,6 +35,7 @@
from tensorflow.python.platform import test
+@test_util.run_all_in_graph_and_eager_modes
def function_set_stats_aggregator(dataset,
aggregator,
prefix="",
@@ -46,7 +46,6 @@
def function_apply_options(dataset, aggregator, prefix="", counter_prefix=""):
options = dataset_ops.Options()
- options.experimental_stats = stats_options.StatsOptions()
options.experimental_stats.aggregator = aggregator
options.experimental_stats.prefix = prefix
options.experimental_stats.counter_prefix = counter_prefix
@@ -60,133 +59,110 @@
)
class StatsDatasetTest(stats_dataset_test_base.StatsDatasetTestBase):
- @test_util.run_deprecated_v1
def testBytesProduced(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).map(
lambda x: array_ops.tile([x], ops.convert_to_tensor([x]))).apply(
stats_ops.bytes_produced_stats("bytes_produced"))
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
+ next_element = self.getNext(dataset, requires_initialization=True)
summary_t = aggregator.get_summary()
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- expected_sum = 0.0
- for i in range(100):
- self.assertAllEqual(
- np.array([i] * i, dtype=np.int64), self.evaluate(next_element))
- summary_str = self.evaluate(summary_t)
- self._assertSummaryHasCount(summary_str, "bytes_produced", float(i + 1))
- expected_sum += i * 8.0
- self._assertSummaryHasSum(summary_str, "bytes_produced", expected_sum)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
- summary_str = self.evaluate(summary_t)
- self._assertSummaryHasCount(summary_str, "bytes_produced", 100.0)
+ expected_sum = 0.0
+ for i in range(100):
+ self.assertAllEqual(
+ np.array([i] * i, dtype=np.int64), self.evaluate(next_element()))
+ summary_str = self.evaluate(aggregator.get_summary())
+ self._assertSummaryHasCount(summary_str, "bytes_produced", float(i + 1))
+ expected_sum += i * 8.0
self._assertSummaryHasSum(summary_str, "bytes_produced", expected_sum)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ # TODO(shivaniagrawal): ntentional breaking case
+ summary_str = self.evaluate(summary_t)
+ self._assertSummaryHasCount(summary_str, "bytes_produced", 100.0)
+ self._assertSummaryHasSum(summary_str, "bytes_produced", expected_sum)
- @test_util.run_deprecated_v1
def testLatencyStats(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency"))
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
+ next_element = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(100):
- self.assertEqual(i, self.evaluate(next_element))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", float(i + 1))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(next_element()))
self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", 100.0)
+ self.evaluate(aggregator.get_summary()), "record_latency",
+ float(i + 1))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency", 100.0)
- @test_util.run_deprecated_v1
def testPrefetchBufferUtilization(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).map(
lambda x: array_ops.tile([x], ops.convert_to_tensor([x]))).prefetch(-1)
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
-
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(100):
- self.assertAllEqual(
- np.array([i] * i, dtype=np.int64), self.evaluate(next_element))
- summary_str = self.evaluate(summary_t)
- self._assertSummaryHasCount(summary_str, "Prefetch::buffer_utilization",
- float(i + 1))
- self._assertSummaryContains(summary_str, "Prefetch::buffer_capacity")
- self._assertSummaryContains(summary_str, "Prefetch::buffer_size")
- self._assertSummaryHasRange(summary_str, "Prefetch::buffer_utilization",
- 0, 1)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
- summary_str = self.evaluate(summary_t)
+ next_element = self.getNext(dataset, requires_initialization=True)
+ for i in range(100):
+ self.assertAllEqual(
+ np.array([i] * i, dtype=np.int64), self.evaluate(next_element()))
+ summary_str = self.evaluate(aggregator.get_summary())
self._assertSummaryHasCount(summary_str, "Prefetch::buffer_utilization",
- 100)
+ float(i + 1))
+ self._assertSummaryContains(summary_str, "Prefetch::buffer_capacity")
+ self._assertSummaryContains(summary_str, "Prefetch::buffer_size")
+ self._assertSummaryHasRange(summary_str, "Prefetch::buffer_utilization",
+ 0, 1)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ summary_str = self.evaluate(aggregator.get_summary())
+ self._assertSummaryHasCount(summary_str, "Prefetch::buffer_utilization",
+ 100)
- @test_util.run_deprecated_v1
def testPrefetchBufferScalars(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(10).map(
- lambda x: array_ops.tile([x], ops.convert_to_tensor([x]))).prefetch(0)
+ lambda x: array_ops.tile([x], ops.convert_to_tensor([x]))).prefetch(1)
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
+ next_element = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(10):
- self.assertAllEqual(
- np.array([i] * i, dtype=np.int64), self.evaluate(next_element))
- summary_str = self.evaluate(summary_t)
- self._assertSummaryHasScalarValue(summary_str,
- "Prefetch::buffer_capacity", 0)
- self._assertSummaryHasScalarValue(summary_str, "Prefetch::buffer_size",
- 0)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(10):
+ self.assertAllEqual(
+ np.array([i] * i, dtype=np.int64), self.evaluate(next_element()))
+ summary_str = self.evaluate(aggregator.get_summary())
+ self._assertSummaryHasScalarValue(summary_str,
+ "Prefetch::buffer_capacity", 1)
+ self._assertSummaryHasScalarValue(summary_str, "Prefetch::buffer_size", 1)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
- @test_util.run_deprecated_v1
def testFilteredElementsStats(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(101).filter(
lambda x: math_ops.equal(math_ops.mod(x, 3), 0))
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
+ next_element = self.getNext(dataset, requires_initialization=True)
- with self.test_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(34):
- self.assertEqual(i * 3, self.evaluate(next_element))
- if i is not 0:
- self._assertSummaryHasScalarValue(
- self.evaluate(summary_t), "Filter::dropped_elements",
- float(i * 2))
- self._assertSummaryHasScalarValue(
- self.evaluate(summary_t), "Filter::filtered_elements", float(i + 1))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(34):
+ self.assertEqual(i * 3, self.evaluate(next_element()))
+ summary_str = self.evaluate(aggregator.get_summary())
+ if i is not 0:
+ self._assertSummaryHasScalarValue(summary_str,
+ "Filter::dropped_elements",
+ float(i * 2))
self._assertSummaryHasScalarValue(
- self.evaluate(summary_t), "Filter::dropped_elements", 67.0)
- self._assertSummaryHasScalarValue(
- self.evaluate(summary_t), "Filter::filtered_elements", 34.0)
+ summary_str, "Filter::filtered_elements", float(i + 1))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ summary_str = self.evaluate(aggregator.get_summary())
+ self._assertSummaryHasScalarValue(summary_str, "Filter::dropped_elements",
+ 67.0)
+ self._assertSummaryHasScalarValue(summary_str, "Filter::filtered_elements",
+ 34.0)
- @test_util.run_deprecated_v1
def testMapBufferUtilization(self, dataset_transformation):
def dataset_fn():
@@ -201,7 +177,6 @@
dataset_transformation,
function_processing_time=True)
- @test_util.run_deprecated_v1
def testMapAutoTuneBufferUtilization(self, dataset_transformation):
def dataset_fn():
@@ -219,7 +194,6 @@
dataset_transformation,
function_processing_time=True)
- @test_util.run_deprecated_v1
def testInterleaveAutoTuneBufferUtilization(self, dataset_transformation):
def dataset_fn():
@@ -236,7 +210,6 @@
self._testParallelCallsStats(dataset_fn, "ParallelInterleaveV2", 10,
dataset_transformation)
- @test_util.run_deprecated_v1
def testMapAndBatchAutoTuneBufferUtilization(self, dataset_transformation):
def dataset_fn():
@@ -258,114 +231,98 @@
check_elements=False,
function_processing_time=True)
- @test_util.run_deprecated_v1
def testReinitialize(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency"))
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
- with self.cached_session() as sess:
- for j in range(5):
- self.evaluate(iterator.initializer)
- for i in range(100):
- self.assertEqual(i, self.evaluate(next_element))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency",
- float((j * 100) + i + 1))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for j in range(5):
+ next_element = self.getNext(dataset, requires_initialization=True)
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(next_element()))
self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", (j + 1) * 100.0)
+ self.evaluate(aggregator.get_summary()), "record_latency",
+ float((j * 100) + i + 1))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency",
+ (j + 1) * 100.0)
- @test_util.run_deprecated_v1
def testNoAggregatorRegistered(self, dataset_transformation):
dataset = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency"))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(100):
- self.assertEqual(i, self.evaluate(next_element))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ next_element = self.getNext(dataset, requires_initialization=True)
- @test_util.run_deprecated_v1
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(next_element()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+
def testMultipleTags(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency")).apply(
stats_ops.latency_stats("record_latency_2"))
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(100):
- self.assertEqual(i, self.evaluate(next_element))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", float(i + 1))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency_2", float(i + 1))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", 100.0)
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency_2", 100.0)
+ next_element = self.getNext(dataset, requires_initialization=True)
- @test_util.run_deprecated_v1
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(next_element()))
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency",
+ float(i + 1))
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency_2",
+ float(i + 1))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency", 100.0)
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency_2", 100.0)
+
def testRepeatedTags(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency")).apply(
stats_ops.latency_stats("record_latency"))
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
+ next_element = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- self.evaluate(iterator.initializer)
- for i in range(100):
- self.assertEqual(i, self.evaluate(next_element))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", float(2 * (i + 1)))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(next_element()))
self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", 200.0)
+ self.evaluate(aggregator.get_summary()), "record_latency",
+ float(2 * (i + 1)))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency", 200.0)
- @test_util.run_deprecated_v1
def testMultipleIteratorsSameAggregator(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency"))
dataset = dataset_transformation(dataset, aggregator)
- iterator_0 = dataset_ops.make_initializable_iterator(dataset)
- iterator_1 = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator_0.get_next() + iterator_1.get_next()
- summary_t = aggregator.get_summary()
+ next_element1 = self.getNext(dataset, requires_initialization=True)
+ next_element2 = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- self.evaluate([iterator_0.initializer, iterator_1.initializer])
- for i in range(100):
- self.assertEqual(i * 2, self.evaluate(next_element))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", float(2 * (i + 1)))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(100):
+ self.assertEqual(i * 2, self.evaluate(next_element1() + next_element2()))
self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_latency", 200.0)
+ self.evaluate(aggregator.get_summary()), "record_latency",
+ float(2 * (i + 1)))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element1())
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element2())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_latency", 200.0)
- @test_util.run_deprecated_v1
def testMultipleDatasetWithPrefixes(self, dataset_transformation):
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_ops.Dataset.range(100).apply(
@@ -374,25 +331,27 @@
dataset2 = dataset_ops.Dataset.range(100).apply(
stats_ops.latency_stats("record_latency"))
dataset2 = dataset_transformation(dataset2, aggregator, prefix="dataset2")
- iterator_0 = dataset_ops.make_initializable_iterator(dataset)
- iterator_1 = dataset_ops.make_initializable_iterator(dataset2)
- next_element = iterator_0.get_next() + iterator_1.get_next()
- summary_t = aggregator.get_summary()
+ next_element1 = self.getNext(dataset, requires_initialization=True)
+ next_element2 = self.getNext(dataset2, requires_initialization=True)
- with self.test_session() as sess:
- self.evaluate([iterator_0.initializer, iterator_1.initializer])
- for i in range(100):
- self.assertEqual(i * 2, self.evaluate(next_element))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "dataset1_record_latency", float(i + 1))
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "dataset2_record_latency", float(i + 1))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(100):
+ self.assertEqual(i * 2, self.evaluate(next_element1() + next_element2()))
self._assertSummaryHasCount(
- self.evaluate(summary_t), "dataset1_record_latency", 100.0)
+ self.evaluate(aggregator.get_summary()), "dataset1_record_latency",
+ float(i + 1))
self._assertSummaryHasCount(
- self.evaluate(summary_t), "dataset2_record_latency", 100.0)
+ self.evaluate(aggregator.get_summary()), "dataset2_record_latency",
+ float(i + 1))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element1())
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element2())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "dataset1_record_latency",
+ 100.0)
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "dataset2_record_latency",
+ 100.0)
@parameterized.named_parameters(
@@ -406,7 +365,6 @@
stats_dataset_test_base.StatsDatasetTestBase,
reader_dataset_ops_test_base.MakeBatchedFeaturesDatasetTestBase):
- @test_util.run_deprecated_v1
def testFeaturesStats(self, dataset_transformation):
num_epochs = 5
total_records = num_epochs * self._num_records
@@ -435,27 +393,26 @@
dataset = dataset_transformation(
dataset_fn(), aggregator, prefix="record_stats")
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
- with self.test_session() as sess:
- self.evaluate(iterator.initializer)
- for _ in range(num_output):
- self.evaluate(next_element)
+ next_element = self.getNext(dataset, requires_initialization=True)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_stats_features", total_records)
- self._assertSummaryHasCount(
- self.evaluate(summary_t), "record_stats_feature-values",
- total_records)
- self._assertSummaryHasSum(
- self.evaluate(summary_t), "record_stats_features", total_records * 4)
- self._assertSummaryHasSum(
- self.evaluate(summary_t), "record_stats_feature-values",
- self._sum_keywords(1) * num_epochs + 3 * total_records)
+ for _ in range(num_output):
+ self.evaluate(next_element())
+
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_stats_features",
+ total_records)
+ self._assertSummaryHasCount(
+ self.evaluate(aggregator.get_summary()), "record_stats_feature-values",
+ total_records)
+ self._assertSummaryHasSum(
+ self.evaluate(aggregator.get_summary()), "record_stats_features",
+ total_records * 4)
+ self._assertSummaryHasSum(
+ self.evaluate(aggregator.get_summary()), "record_stats_feature-values",
+ self._sum_keywords(1) * num_epochs + 3 * total_records)
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py
index ab1d1c3..b80aab9 100644
--- a/tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py
+++ b/tensorflow/python/data/experimental/kernel_tests/stats_dataset_test_base.py
@@ -22,7 +22,6 @@
from tensorflow.core.framework import summary_pb2
from tensorflow.python.data.experimental.ops import stats_aggregator
from tensorflow.python.data.kernel_tests import test_base
-from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import errors
@@ -94,27 +93,23 @@
aggregator = stats_aggregator.StatsAggregator()
dataset = dataset_fn()
dataset = dataset_transformation(dataset, aggregator)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- summary_t = aggregator.get_summary()
+ next_element = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- sess.run(iterator.initializer)
- for i in range(num_output):
- next_ = sess.run(next_element)
- if check_elements:
- self.assertAllEqual(np.array([i] * i, dtype=np.int64), next_)
- summary_str = sess.run(summary_t)
- if function_processing_time:
- self._assertSummaryHasCountMoreOrEqualGeneralisedTag(
- summary_str, "::execution_time", float(i + 1))
- self._assertSummaryContains(summary_str,
- dataset_name + "::num_parallel_calls")
- self._assertSummaryContains(summary_str,
- dataset_name + "::active_parallel_calls")
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_element)
+ for i in range(num_output):
+ next_ = self.evaluate(next_element())
+ if check_elements:
+ self.assertAllEqual(np.array([i] * i, dtype=np.int64), next_)
+ summary_str = self.evaluate(aggregator.get_summary())
if function_processing_time:
- summary_str = sess.run(summary_t)
self._assertSummaryHasCountMoreOrEqualGeneralisedTag(
- summary_str, "::execution_time", float(num_output))
+ summary_str, "::execution_time", float(i + 1))
+ self._assertSummaryContains(summary_str,
+ dataset_name + "::num_parallel_calls")
+ self._assertSummaryContains(summary_str,
+ dataset_name + "::active_parallel_calls")
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
+ if function_processing_time:
+ summary_str = self.evaluate(aggregator.get_summary())
+ self._assertSummaryHasCountMoreOrEqualGeneralisedTag(
+ summary_str, "::execution_time", float(num_output))
diff --git a/tensorflow/python/data/experimental/kernel_tests/tf_record_writer_test.py b/tensorflow/python/data/experimental/kernel_tests/tf_record_writer_test.py
index 8fd0ad5..14a4241 100644
--- a/tensorflow/python/data/experimental/kernel_tests/tf_record_writer_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/tf_record_writer_test.py
@@ -23,26 +23,24 @@
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import readers
-from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
from tensorflow.python.lib.io import python_io
from tensorflow.python.lib.io import tf_record
-from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
from tensorflow.python.util import compat
+@test_util.run_all_in_graph_and_eager_modes
class TFRecordWriterTest(test_base.DatasetTestBase):
def setUp(self):
super(TFRecordWriterTest, self).setUp()
self._num_records = 7
- self.filename = array_ops.placeholder(dtypes.string, shape=[])
- self.compression_type = array_ops.placeholder_with_default("", shape=[])
- input_dataset = readers.TFRecordDataset([self.filename],
- self.compression_type)
- self.writer = writers.TFRecordWriter(
- self._outputFilename(), self.compression_type).write(input_dataset)
+ def writer_fn(self, filename, compression_type=""):
+ input_dataset = readers.TFRecordDataset([filename], compression_type)
+ return writers.TFRecordWriter(self._outputFilename(),
+ compression_type).write(input_dataset)
def _record(self, i):
return compat.as_bytes("Record %d" % (i))
@@ -62,56 +60,39 @@
return os.path.join(self.get_temp_dir(), "tf_record.out.txt")
def testWrite(self):
- with self.cached_session() as sess:
- sess.run(
- self.writer, feed_dict={
- self.filename: self._createFile(),
- })
+ self.evaluate(self.writer_fn(self._createFile()))
for i, r in enumerate(tf_record.tf_record_iterator(self._outputFilename())):
self.assertAllEqual(self._record(i), r)
def testWriteZLIB(self):
options = tf_record.TFRecordOptions(tf_record.TFRecordCompressionType.ZLIB)
- with self.cached_session() as sess:
- sess.run(
- self.writer,
- feed_dict={
- self.filename: self._createFile(options),
- self.compression_type: "ZLIB",
- })
+ self.evaluate(
+ self.writer_fn(self._createFile(options), compression_type="ZLIB"))
for i, r in enumerate(
tf_record.tf_record_iterator(self._outputFilename(), options=options)):
self.assertAllEqual(self._record(i), r)
def testWriteGZIP(self):
options = tf_record.TFRecordOptions(tf_record.TFRecordCompressionType.GZIP)
- with self.cached_session() as sess:
- sess.run(
- self.writer,
- feed_dict={
- self.filename: self._createFile(options),
- self.compression_type: "GZIP",
- })
+ self.evaluate(
+ self.writer_fn(self._createFile(options), compression_type="GZIP"))
for i, r in enumerate(
tf_record.tf_record_iterator(self._outputFilename(), options=options)):
self.assertAllEqual(self._record(i), r)
def testFailDataset(self):
with self.assertRaises(TypeError):
- writers.TFRecordWriter(self._outputFilename(),
- self.compression_type).write("whoops")
+ writers.TFRecordWriter(self._outputFilename(), "").write("whoops")
def testFailDType(self):
input_dataset = dataset_ops.Dataset.from_tensors(10)
with self.assertRaises(TypeError):
- writers.TFRecordWriter(self._outputFilename(),
- self.compression_type).write(input_dataset)
+ writers.TFRecordWriter(self._outputFilename(), "").write(input_dataset)
def testFailShape(self):
input_dataset = dataset_ops.Dataset.from_tensors([["hello"], ["world"]])
with self.assertRaises(TypeError):
- writers.TFRecordWriter(self._outputFilename(),
- self.compression_type).write(input_dataset)
+ writers.TFRecordWriter(self._outputFilename(), "").write(input_dataset)
if __name__ == "__main__":
diff --git a/tensorflow/python/data/experimental/kernel_tests/unbatch_test.py b/tensorflow/python/data/experimental/kernel_tests/unbatch_test.py
index cef5e8d..7f30cf5 100644
--- a/tensorflow/python/data/experimental/kernel_tests/unbatch_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/unbatch_test.py
@@ -36,24 +36,14 @@
from tensorflow.python.util import compat
+@test_util.run_all_in_graph_and_eager_modes
class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase):
- @test_util.run_deprecated_v1
def testUnbatchWithUnknownRankInput(self):
- placeholder = array_ops.placeholder(dtypes.int32)
- dataset = dataset_ops.Dataset.from_tensors(placeholder).apply(
- batching.unbatch())
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_elem = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors([0, 1, 2,
+ 3]).apply(batching.unbatch())
+ self.assertDatasetProduces(dataset, range(4))
- with self.cached_session() as sess:
- sess.run(iterator.initializer, feed_dict={placeholder: [0, 1, 2, 3]})
- for i in range(4):
- self.assertEqual(i, self.evaluate(next_elem))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_elem)
-
- @test_util.run_deprecated_v1
def testUnbatchScalarDataset(self):
data = tuple([math_ops.range(10) for _ in range(3)])
data = dataset_ops.Dataset.from_tensor_slices(data)
@@ -63,17 +53,8 @@
data = data.apply(batching.unbatch())
self.assertEqual(expected_types, data.output_types)
- iterator = dataset_ops.make_one_shot_iterator(data)
- op = iterator.get_next()
+ self.assertDatasetProduces(data, [(i,) * 3 for i in range(10)])
- with self.cached_session() as sess:
- for i in range(10):
- self.assertEqual((i,) * 3, self.evaluate(op))
-
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(op)
-
- @test_util.run_deprecated_v1
def testUnbatchDatasetWithStrings(self):
data = tuple([math_ops.range(10) for _ in range(3)])
data = dataset_ops.Dataset.from_tensor_slices(data)
@@ -84,18 +65,12 @@
data = data.apply(batching.unbatch())
self.assertEqual(expected_types, data.output_types)
- iterator = dataset_ops.make_one_shot_iterator(data)
- op = iterator.get_next()
+ self.assertDatasetProduces(
+ data, [(i, compat.as_bytes(str(i)), i) for i in range(10)])
- with self.cached_session() as sess:
- for i in range(10):
- self.assertEqual((i, compat.as_bytes(str(i)), i), self.evaluate(op))
-
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(op)
-
+ # TODO(b/119837791): Add eager coverage.
@test_util.run_deprecated_v1
- def testUnbatchDatasetWithSparseTensor(self):
+ def testSkipEagerUnbatchDatasetWithSparseTensor(self):
st = sparse_tensor.SparseTensorValue(
indices=[[i, i] for i in range(10)],
values=list(range(10)),
@@ -107,17 +82,17 @@
iterator = dataset_ops.make_one_shot_iterator(data)
next_element = iterator.get_next()
- with self.cached_session() as sess:
- for i in range(10):
- st_row = self.evaluate(next_element)
- self.assertEqual([i], st_row.indices)
- self.assertEqual([i], st_row.values)
- self.assertEqual([10], st_row.dense_shape)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(10):
+ st_row = self.evaluate(next_element)
+ self.assertEqual([i], st_row.indices)
+ self.assertEqual([i], st_row.values)
+ self.assertEqual([10], st_row.dense_shape)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element)
+ # TODO(b/119837791): Add eager coverage.
@test_util.run_deprecated_v1
- def testUnbatchDatasetWithDenseAndSparseTensor(self):
+ def testSkipEagerUnbatchDatasetWithDenseAndSparseTensor(self):
st = sparse_tensor.SparseTensorValue(
indices=[[i, i] for i in range(10)],
values=list(range(10)),
@@ -126,20 +101,17 @@
data = data.apply(batching.unbatch())
data = data.batch(5)
data = data.apply(batching.unbatch())
- iterator = dataset_ops.make_one_shot_iterator(data)
- next_element = iterator.get_next()
+ next_element = self.getNext(data)
- with self.cached_session() as sess:
- for i in range(10):
- dense_elem, st_row = self.evaluate(next_element)
- self.assertEqual(i, dense_elem)
- self.assertEqual([i], st_row.indices)
- self.assertEqual([i], st_row.values)
- self.assertEqual([10], st_row.dense_shape)
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for i in range(10):
+ dense_elem, st_row = self.evaluate(next_element())
+ self.assertEqual(i, dense_elem)
+ self.assertEqual([i], st_row.indices)
+ self.assertEqual([i], st_row.values)
+ self.assertEqual([10], st_row.dense_shape)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_element())
- @test_util.run_deprecated_v1
def testUnbatchSingleElementTupleDataset(self):
data = tuple([(math_ops.range(10),) for _ in range(3)])
data = dataset_ops.Dataset.from_tensor_slices(data)
@@ -149,17 +121,8 @@
data = data.apply(batching.unbatch())
self.assertEqual(expected_types, data.output_types)
- iterator = dataset_ops.make_one_shot_iterator(data)
- op = iterator.get_next()
+ self.assertDatasetProduces(data, [((i,),) * 3 for i in range(10)])
- with self.cached_session() as sess:
- for i in range(10):
- self.assertEqual(((i,),) * 3, self.evaluate(op))
-
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(op)
-
- @test_util.run_deprecated_v1
def testUnbatchMultiElementTupleDataset(self):
data = tuple([(math_ops.range(10 * i, 10 * i + 10),
array_ops.fill([10], "hi")) for i in range(3)])
@@ -170,29 +133,16 @@
data = data.apply(batching.unbatch())
self.assertAllEqual(expected_types, data.output_types)
- iterator = dataset_ops.make_one_shot_iterator(data)
- op = iterator.get_next()
+ self.assertDatasetProduces(
+ data,
+ [((i, b"hi"), (10 + i, b"hi"), (20 + i, b"hi")) for i in range(10)])
- with self.cached_session() as sess:
- for i in range(10):
- self.assertEqual(((i, b"hi"), (10 + i, b"hi"), (20 + i, b"hi")),
- self.evaluate(op))
-
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(op)
-
- @test_util.run_deprecated_v1
def testUnbatchEmpty(self):
data = dataset_ops.Dataset.from_tensors(
(constant_op.constant([]), constant_op.constant([], shape=[0, 4]),
constant_op.constant([], shape=[0, 4, 0])))
data = data.apply(batching.unbatch())
- iterator = dataset_ops.make_one_shot_iterator(data)
- next_element = iterator.get_next()
-
- with self.cached_session() as sess:
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ self.assertDatasetProduces(data, [])
def testUnbatchStaticShapeMismatch(self):
data = dataset_ops.Dataset.from_tensors((np.arange(7), np.arange(8),
@@ -200,8 +150,9 @@
with self.assertRaises(ValueError):
data.apply(batching.unbatch())
+ # TODO(b/119837791): eager mode doesnt capture raised error, debug.
@test_util.run_deprecated_v1
- def testUnbatchDynamicShapeMismatch(self):
+ def testSkipEagerUnbatchDynamicShapeMismatch(self):
ph1 = array_ops.placeholder(dtypes.int32, shape=[None])
ph2 = array_ops.placeholder(dtypes.int32, shape=None)
data = dataset_ops.Dataset.from_tensors((ph1, ph2))
diff --git a/tensorflow/python/data/experimental/kernel_tests/unique_test.py b/tensorflow/python/data/experimental/kernel_tests/unique_test.py
index 1d9941d..42d76a2 100644
--- a/tensorflow/python/data/experimental/kernel_tests/unique_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/unique_test.py
@@ -21,12 +21,12 @@
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.util import compat
+@test_util.run_all_in_graph_and_eager_modes
class UniqueTest(test_base.DatasetTestBase):
def _testSimpleHelper(self, dtype, test_cases):
@@ -44,19 +44,13 @@
current_test_case = []
dataset = dataset_ops.Dataset.from_generator(lambda: current_test_case,
dtype).apply(unique.unique())
- iterator = dataset_ops.make_initializable_iterator(dataset)
- next_element = iterator.get_next()
- with self.cached_session() as sess:
- for test_case, expected in test_cases:
- current_test_case = test_case
- self.evaluate(iterator.initializer)
- for element in expected:
- if dtype == dtypes.string:
- element = compat.as_bytes(element)
- self.assertAllEqual(element, self.evaluate(next_element))
- with self.assertRaises(errors.OutOfRangeError):
- self.evaluate(next_element)
+ for test_case, expected in test_cases:
+ current_test_case = test_case
+ self.assertDatasetProduces(dataset, [
+ compat.as_bytes(element) if dtype == dtypes.string else element
+ for element in expected
+ ])
@test_util.run_deprecated_v1
def testSimpleInt(self):
diff --git a/tensorflow/python/data/experimental/kernel_tests/wrap_unwrap_test.py b/tensorflow/python/data/experimental/kernel_tests/wrap_unwrap_test.py
new file mode 100644
index 0000000..c47595c
--- /dev/null
+++ b/tensorflow/python/data/experimental/kernel_tests/wrap_unwrap_test.py
@@ -0,0 +1,68 @@
+# 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 Wrapping / Unwrapping dataset variants."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.platform import test
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class WrapDatasetVariantTest(test_base.DatasetTestBase):
+
+ def testBasic(self):
+ ds = dataset_ops.Dataset.range(100)
+ ds_variant = ds._as_variant_tensor() # pylint: disable=protected-access
+
+ wrapped_variant = gen_dataset_ops.wrap_dataset_variant(ds_variant)
+ unwrapped_variant = gen_dataset_ops.unwrap_dataset_variant(wrapped_variant)
+
+ variant_ds = dataset_ops._VariantDataset(unwrapped_variant,
+ ds._element_structure)
+ get_next = self.getNext(variant_ds, requires_initialization=True)
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(get_next()))
+
+ # TODO(b/119837791): add eager coverage when supported.
+ def testSkipEagerGPU(self):
+ ds = dataset_ops.Dataset.range(100)
+ ds_variant = ds._as_variant_tensor() # pylint: disable=protected-access
+ wrapped_variant = gen_dataset_ops.wrap_dataset_variant(ds_variant)
+
+ with ops.device("/gpu:0"):
+ gpu_wrapped_variant = array_ops.identity(wrapped_variant)
+
+ unwrapped_variant = gen_dataset_ops.unwrap_dataset_variant(
+ gpu_wrapped_variant)
+ variant_ds = dataset_ops._VariantDataset(unwrapped_variant,
+ ds._element_structure)
+ iterator = dataset_ops.make_initializable_iterator(variant_ds)
+ get_next = iterator.get_next()
+
+ with self.cached_session():
+ self.evaluate(iterator.initializer)
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(get_next))
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/data/experimental/ops/BUILD b/tensorflow/python/data/experimental/ops/BUILD
index 43f4318..60c20e0b 100644
--- a/tensorflow/python/data/experimental/ops/BUILD
+++ b/tensorflow/python/data/experimental/ops/BUILD
@@ -70,8 +70,7 @@
"//tensorflow/python:random_seed",
"//tensorflow/python:tensor_shape",
"//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
- "//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/util:structure",
],
)
@@ -230,6 +229,8 @@
"//tensorflow/python:dataset_ops_gen",
"//tensorflow/python:framework_ops",
"//tensorflow/python:tensor_shape",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/util:structure",
],
)
@@ -268,7 +269,7 @@
"//tensorflow/python:sparse_tensor",
"//tensorflow/python:tensor_shape",
"//tensorflow/python/data/ops:dataset_ops",
- "//tensorflow/python/data/util:nest",
+ "//tensorflow/python/data/util:structure",
],
)
@@ -303,6 +304,7 @@
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/util:nest",
"//tensorflow/python/data/util:sparse",
+ "//tensorflow/python/data/util:structure",
],
)
diff --git a/tensorflow/python/data/experimental/ops/batching.py b/tensorflow/python/data/experimental/ops/batching.py
index 668bf3e..29df98f 100644
--- a/tensorflow/python/data/experimental/ops/batching.py
+++ b/tensorflow/python/data/experimental/ops/batching.py
@@ -24,7 +24,7 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import convert
from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
@@ -365,23 +365,19 @@
"different batch sizes.")
self._input_dataset = input_dataset
+ self._structure = structure.convert_legacy_structure(
+ input_dataset.output_types,
+ nest.map_structure(lambda s: s[1:], input_dataset.output_shapes),
+ input_dataset.output_classes)
+
def _as_variant_tensor(self):
return ged_ops.experimental_unbatch_dataset(
self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
**dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return self._input_dataset.output_classes
-
- @property
- def output_shapes(self):
- return nest.map_structure(lambda s: s[1:],
- self._input_dataset.output_shapes)
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
+ def _element_structure(self):
+ return self._structure
@tf_export("data.experimental.unbatch")
@@ -409,21 +405,19 @@
def _apply_fn(dataset):
"""Function from `Dataset` to `Dataset` that applies the transformation."""
- if not sparse.any_sparse(dataset.output_classes):
- return _UnbatchDataset(dataset)
-
# NOTE(mrry): We must ensure that any SparseTensors in `dataset`
# are normalized to the rank-1 dense representation, so that the
# sparse-oblivious unbatching logic will slice them
# appropriately. This leads to a somewhat inefficient re-encoding step
# for all SparseTensor components.
- # TODO(mrry): Consider optimizing this in future
- # if it turns out to be a bottleneck.
+ # TODO(mrry): Consider optimizing this in future if it turns out to be
+ # a bottleneck.
def normalize(arg, *rest):
+ # pylint: disable=protected-access
if rest:
- return sparse.serialize_many_sparse_tensors((arg,) + rest)
+ return dataset._element_structure._to_batched_tensor_list((arg,) + rest)
else:
- return sparse.serialize_many_sparse_tensors(arg)
+ return dataset._element_structure._to_batched_tensor_list(arg)
normalized_dataset = dataset.map(normalize)
@@ -454,6 +448,9 @@
self._input_dataset = input_dataset
self._batch_size = batch_size
self._row_shape = row_shape
+ self._structure = structure.SparseTensorStructure(
+ input_dataset.output_types,
+ tensor_shape.vector(None).concatenate(self._row_shape))
def _as_variant_tensor(self):
return ged_ops.experimental_dense_to_sparse_batch_dataset(
@@ -463,16 +460,8 @@
**dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return sparse_tensor.SparseTensor
-
- @property
- def output_shapes(self):
- return tensor_shape.vector(None).concatenate(self._row_shape)
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
+ def _element_structure(self):
+ return self._structure
class _RestructuredDataset(dataset_ops.UnaryDataset):
@@ -523,13 +512,10 @@
"Dataset with output types %r cannot be restructured to have "
"output types %r" % (dataset.output_types, output_types))
- self._output_types = output_types
-
if output_shapes is None:
# Inherit shapes from the original `dataset`.
- self._output_shapes = nest.pack_sequence_as(output_types,
- nest.flatten(
- dataset.output_shapes))
+ output_shapes = nest.pack_sequence_as(
+ output_types, nest.flatten(dataset.output_shapes))
else:
if not allow_unsafe_cast:
# Validate that the shapes are compatible.
@@ -544,30 +530,22 @@
"Dataset with output shapes %r cannot be restructured to have "
"incompatible output shapes %r" % (dataset.output_shapes,
output_shapes))
- self._output_shapes = nest.map_structure_up_to(
+ output_shapes = nest.map_structure_up_to(
output_types, tensor_shape.as_shape, output_shapes)
if output_classes is None:
# Inherit class types from the original `dataset`.
- self._output_classes = nest.pack_sequence_as(output_types,
- nest.flatten(
- dataset.output_classes))
- else:
- self._output_classes = output_classes
+ output_classes = nest.pack_sequence_as(
+ output_types, nest.flatten(dataset.output_classes))
+
+ self._structure = structure.convert_legacy_structure(
+ output_types, output_shapes, output_classes)
def _as_variant_tensor(self):
return self._input_dataset._as_variant_tensor() # pylint: disable=protected-access
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_types(self):
- return self._output_types
-
- @property
- def output_shapes(self):
- return self._output_shapes
+ def _element_structure(self):
+ return self._structure
class _MapAndBatchDataset(dataset_ops.UnaryDataset):
@@ -591,10 +569,10 @@
if constant_drop_remainder:
# NOTE(mrry): `constant_drop_remainder` may be `None` (unknown statically)
# or `False` (explicitly retaining the remainder).
- self._output_structure = self._map_func.output_structure._batch( # pylint: disable=protected-access
+ self._structure = self._map_func.output_structure._batch( # pylint: disable=protected-access
tensor_util.constant_value(self._batch_size_t))
else:
- self._output_structure = self._map_func.output_structure._batch(None) # pylint: disable=protected-access
+ self._structure = self._map_func.output_structure._batch(None) # pylint: disable=protected-access
def _functions(self):
return [self._map_func]
@@ -608,19 +586,12 @@
batch_size=self._batch_size_t,
num_parallel_calls=self._num_parallel_calls_t,
drop_remainder=self._drop_remainder_t,
- **dataset_ops.flat_structure(structure=self._output_structure))
+ preserve_cardinality=True,
+ **dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return self._output_structure._to_legacy_output_classes() # pylint: disable=protected-access
-
- @property
- def output_shapes(self):
- return self._output_structure._to_legacy_output_shapes() # pylint: disable=protected-access
-
- @property
- def output_types(self):
- return self._output_structure._to_legacy_output_types() # pylint: disable=protected-access
+ def _element_structure(self):
+ return self._structure
@tf_export("data.experimental.map_and_batch")
@@ -652,9 +623,10 @@
whether the last batch should be dropped in case its size is smaller than
desired; the default behavior is not to drop the smaller batch.
num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
- representing the number of elements to process in parallel. If not
- specified, `batch_size * num_parallel_batches` elements will be
- processed in parallel.
+ representing the number of elements to process in parallel. If not
+ specified, `batch_size * num_parallel_batches` elements will be processed
+ in parallel. If the value `tf.data.experimental.AUTOTUNE` is used, then
+ the number of parallel calls is set dynamically based on available CPU.
Returns:
A `Dataset` transformation function, which can be passed to
diff --git a/tensorflow/python/data/experimental/ops/get_single_element.py b/tensorflow/python/data/experimental/ops/get_single_element.py
index 73116ed..d649a07 100644
--- a/tensorflow/python/data/experimental/ops/get_single_element.py
+++ b/tensorflow/python/data/experimental/ops/get_single_element.py
@@ -18,8 +18,6 @@
from __future__ import print_function
from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
from tensorflow.python.ops import gen_dataset_ops
from tensorflow.python.util.tf_export import tf_export
@@ -63,10 +61,8 @@
if not isinstance(dataset, dataset_ops.DatasetV2):
raise TypeError("`dataset` must be a `tf.data.Dataset` object.")
- nested_ret = nest.pack_sequence_as(
- dataset.output_types, gen_dataset_ops.dataset_to_single_element(
- dataset._as_variant_tensor(), # pylint: disable=protected-access
+ # pylint: disable=protected-access
+ return dataset._element_structure._from_compatible_tensor_list(
+ gen_dataset_ops.dataset_to_single_element(
+ dataset._as_variant_tensor(),
**dataset_ops.flat_structure(dataset)))
- return sparse.deserialize_sparse_tensors(
- nested_ret, dataset.output_types, dataset.output_shapes,
- dataset.output_classes)
diff --git a/tensorflow/python/data/experimental/ops/grouping.py b/tensorflow/python/data/experimental/ops/grouping.py
index ad9fe9a..ef6b232 100644
--- a/tensorflow/python/data/experimental/ops/grouping.py
+++ b/tensorflow/python/data/experimental/ops/grouping.py
@@ -253,50 +253,44 @@
def _make_key_func(self, key_func, input_dataset):
"""Make wrapping defun for key_func."""
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ self._key_func = dataset_ops.StructuredFunctionWrapper(
key_func, self._transformation_name(), dataset=input_dataset)
- if not (
- wrapped_func.output_types == dtypes.int64 and
- wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
+ if not self._key_func.output_structure.is_compatible_with(
+ structure.TensorStructure(dtypes.int64, [])):
raise ValueError(
"`key_func` must return a single tf.int64 tensor. "
"Got type=%s and shape=%s"
- % (wrapped_func.output_types, wrapped_func.output_shapes))
- self._key_func = wrapped_func
-
+ % (self._key_func.output_types, self._key_func.output_shapes))
def _make_init_func(self, init_func):
"""Make wrapping defun for init_func."""
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ self._init_func = dataset_ops.StructuredFunctionWrapper(
init_func,
self._transformation_name(),
- input_classes=ops.Tensor,
- input_shapes=tensor_shape.scalar(),
- input_types=dtypes.int64)
- self._init_func = wrapped_func
- self._state_classes = wrapped_func.output_classes
- self._state_shapes = wrapped_func.output_shapes
- self._state_types = wrapped_func.output_types
+ input_structure=structure.TensorStructure(dtypes.int64, []))
def _make_reduce_func(self, reduce_func, input_dataset):
"""Make wrapping defun for reduce_func."""
# Iteratively rerun the reduce function until reaching a fixed point on
- # `self._state_shapes`.
+ # `self._state_structure`.
+ self._state_structure = self._init_func.output_structure
+ state_types = self._init_func.output_types
+ state_shapes = self._init_func.output_shapes
+ state_classes = self._init_func.output_classes
need_to_rerun = True
while need_to_rerun:
wrapped_func = dataset_ops.StructuredFunctionWrapper(
reduce_func,
self._transformation_name(),
- input_classes=(self._state_classes, input_dataset.output_classes),
- input_shapes=(self._state_shapes, input_dataset.output_shapes),
- input_types=(self._state_types, input_dataset.output_types),
+ input_structure=structure.NestedStructure(
+ (self._state_structure, input_dataset._element_structure)), # pylint: disable=protected-access
add_to_graph=False)
# Extract and validate class information from the returned values.
for new_state_class, state_class in zip(
nest.flatten(wrapped_func.output_classes),
- nest.flatten(self._state_classes)):
+ nest.flatten(state_classes)):
if not issubclass(new_state_class, state_class):
raise TypeError(
"The element classes for the new state must match the initial "
@@ -305,16 +299,15 @@
# Extract and validate type information from the returned values.
for new_state_type, state_type in zip(
- nest.flatten(wrapped_func.output_types),
- nest.flatten(self._state_types)):
+ nest.flatten(wrapped_func.output_types), nest.flatten(state_types)):
if new_state_type != state_type:
raise TypeError(
"The element types for the new state must match the initial "
"state. Expected %s; got %s." %
- (self._state_types, wrapped_func.output_types))
+ (self._init_func.output_types, wrapped_func.output_types))
# Extract shape information from the returned values.
- flat_state_shapes = nest.flatten(self._state_shapes)
+ flat_state_shapes = nest.flatten(state_shapes)
flat_new_state_shapes = nest.flatten(wrapped_func.output_shapes)
weakened_state_shapes = [
original.most_specific_compatible_shape(new)
@@ -331,36 +324,23 @@
break
if need_to_rerun:
- self._state_shapes = nest.pack_sequence_as(self._state_shapes,
- weakened_state_shapes)
+ state_shapes = nest.pack_sequence_as(
+ self._init_func.output_shapes, weakened_state_shapes)
+ self._state_structure = structure.convert_legacy_structure(
+ state_types, state_shapes, state_classes)
self._reduce_func = wrapped_func
self._reduce_func.function.add_to_graph(ops.get_default_graph())
def _make_finalize_func(self, finalize_func):
"""Make wrapping defun for finalize_func."""
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- finalize_func,
- self._transformation_name(),
- input_classes=self._state_classes,
- input_shapes=self._state_shapes,
- input_types=self._state_types)
- self._finalize_func = wrapped_func
- self._output_classes = wrapped_func.output_classes
- self._output_shapes = wrapped_func.output_shapes
- self._output_types = wrapped_func.output_types
+ self._finalize_func = dataset_ops.StructuredFunctionWrapper(
+ finalize_func, self._transformation_name(),
+ input_structure=self._state_structure)
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._finalize_func.output_structure
def _functions(self):
return [
@@ -402,66 +382,46 @@
def window_size_func_wrapper(key):
return ops.convert_to_tensor(window_size_func(key), dtype=dtypes.int64)
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ self._window_size_func = dataset_ops.StructuredFunctionWrapper(
window_size_func_wrapper,
self._transformation_name(),
- input_classes=ops.Tensor,
- input_shapes=tensor_shape.scalar(),
- input_types=dtypes.int64)
- if not (
- wrapped_func.output_types == dtypes.int64 and
- wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
+ input_structure=structure.TensorStructure(dtypes.int64, []))
+ if not self._window_size_func.output_structure.is_compatible_with(
+ structure.TensorStructure(dtypes.int64, [])):
raise ValueError(
"`window_size_func` must return a single tf.int64 scalar tensor.")
- self._window_size_func = wrapped_func
def _make_key_func(self, key_func, input_dataset):
"""Make wrapping defun for key_func."""
def key_func_wrapper(*args):
return ops.convert_to_tensor(key_func(*args), dtype=dtypes.int64)
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
+ self._key_func = dataset_ops.StructuredFunctionWrapper(
key_func_wrapper, self._transformation_name(), dataset=input_dataset)
- if not (
- wrapped_func.output_types == dtypes.int64 and
- wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
+ if not self._key_func.output_structure.is_compatible_with(
+ structure.TensorStructure(dtypes.int64, [])):
raise ValueError(
"`key_func` must return a single tf.int64 scalar tensor.")
- self._key_func = wrapped_func
def _make_reduce_func(self, reduce_func, input_dataset):
"""Make wrapping defun for reduce_func."""
nested_dataset = dataset_ops.DatasetStructure(
- structure.Structure._from_legacy_structure( # pylint: disable=protected-access
- input_dataset.output_types, input_dataset.output_shapes,
- input_dataset.output_classes))
- wrapped_func = dataset_ops.StructuredFunctionWrapper(
- reduce_func,
- self._transformation_name(),
- input_classes=(ops.Tensor, nested_dataset),
- input_shapes=(tensor_shape.scalar(), nested_dataset),
- input_types=(dtypes.int64, nested_dataset))
+ input_dataset._element_structure) # pylint: disable=protected-access
+ input_structure = structure.NestedStructure(
+ (structure.TensorStructure(dtypes.int64, []), nested_dataset))
+ self._reduce_func = dataset_ops.StructuredFunctionWrapper(
+ reduce_func, self._transformation_name(),
+ input_structure=input_structure)
if not isinstance(
- wrapped_func.output_structure, dataset_ops.DatasetStructure):
+ self._reduce_func.output_structure, dataset_ops.DatasetStructure):
raise TypeError("`reduce_func` must return a `Dataset` object.")
# pylint: disable=protected-access
- element_structure = wrapped_func.output_structure._element_structure
- self._output_classes = element_structure._to_legacy_output_classes()
- self._output_types = element_structure._to_legacy_output_types()
- self._output_shapes = element_structure._to_legacy_output_shapes()
- self._reduce_func = wrapped_func
+ self._structure = (
+ self._reduce_func.output_structure._element_structure)
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
def _functions(self):
return [self._key_func, self._reduce_func, self._window_size_func]
diff --git a/tensorflow/python/data/experimental/ops/indexed_dataset_ops.py b/tensorflow/python/data/experimental/ops/indexed_dataset_ops.py
index 570f011..fdf3692 100644
--- a/tensorflow/python/data/experimental/ops/indexed_dataset_ops.py
+++ b/tensorflow/python/data/experimental/ops/indexed_dataset_ops.py
@@ -22,9 +22,9 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
from tensorflow.python.data.util import sparse
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
@@ -94,11 +94,7 @@
ged_ops.experimental_materialized_index_dataset_handle(
container=container,
shared_name=shared_name,
- output_types=nest.flatten(
- sparse.as_dense_types(self.output_types, self.output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_types(self.output_shapes,
- self.output_classes))))
+ **dataset_ops.flat_structure(self)))
with ops.colocate_with(materialized_resource):
materializer = ged_ops.experimental_indexed_dataset_materialize(
@@ -107,38 +103,6 @@
self.output_classes, self.output_types,
self.output_shapes)
- @abc.abstractproperty
- def output_types(self):
- """Returns the type of each component of an element of this IndexedDataset.
-
- Returns:
- A nested structure of `tf.DType` objects corresponding to each component
- of an element of this IndexedDataset.
- """
- raise NotImplementedError("IndexedDataset.output_types")
-
- @abc.abstractproperty
- def output_classes(self):
- """Returns the class of each component of an element of this IndexedDataset.
-
- The expected values are `tf.Tensor` and `tf.SparseTensor`.
-
- Returns:
- A nested structure of Python `type` objects corresponding to each
- component of an element of this IndexedDataset.
- """
- raise NotImplementedError("IndexedDataset.output_classes")
-
- @abc.abstractproperty
- def output_shapes(self):
- """Returns the shape of each component of an element of this IndexedDataset.
-
- Returns:
- A nested structure of `tf.TensorShape` objects corresponding to each
- component of an element of this IndexedDataset.
- """
- raise NotImplementedError("IndexedDataset.output_shapes")
-
@abc.abstractmethod
def _as_variant_tensor(self):
"""Creates a `tf.variant` `tf.Tensor` representing this IndexedDataset.
@@ -161,16 +125,8 @@
self._size = ops.convert_to_tensor(size, dtype=dtypes.uint64, name="size")
@property
- def output_types(self):
- return dtypes.uint64
-
- @property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.uint64, [])
def _as_variant_tensor(self):
return ged_ops.experimental_identity_indexed_dataset(self._size)
diff --git a/tensorflow/python/data/experimental/ops/interleave_ops.py b/tensorflow/python/data/experimental/ops/interleave_ops.py
index 8b0fdfc..5a719f8 100644
--- a/tensorflow/python/data/experimental/ops/interleave_ops.py
+++ b/tensorflow/python/data/experimental/ops/interleave_ops.py
@@ -21,6 +21,7 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import readers
from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
@@ -101,6 +102,18 @@
data_input.output_classes != data_inputs[0].output_classes):
raise TypeError("All datasets must have the same type and class.")
+ output_shapes = self._data_inputs[0].output_shapes
+ for data_input in self._data_inputs[1:]:
+ output_shapes = nest.pack_sequence_as(output_shapes, [
+ ts1.most_specific_compatible_shape(ts2) for (ts1, ts2) in zip(
+ nest.flatten(output_shapes),
+ nest.flatten(data_input.output_shapes))
+ ])
+
+ self._structure = structure.convert_legacy_structure(
+ data_inputs[0].output_types, output_shapes,
+ data_inputs[0].output_classes)
+
def _as_variant_tensor(self):
# pylint: disable=protected-access
return (
@@ -115,22 +128,8 @@
return [self._selector_input] + self._data_inputs
@property
- def output_classes(self):
- return self._data_inputs[0].output_classes
-
- @property
- def output_shapes(self):
- ret = self._data_inputs[0].output_shapes
- for data_input in self._data_inputs[1:]:
- ret = nest.pack_sequence_as(ret, [
- ts1.most_specific_compatible_shape(ts2) for (ts1, ts2) in zip(
- nest.flatten(ret), nest.flatten(data_input.output_shapes))
- ])
- return ret
-
- @property
- def output_types(self):
- return self._data_inputs[0].output_types
+ def _element_structure(self):
+ return self._structure
@tf_export("data.experimental.sample_from_datasets", v1=[])
diff --git a/tensorflow/python/data/experimental/ops/matching_files.py b/tensorflow/python/data/experimental/ops/matching_files.py
index 8398f86..63b99cb 100644
--- a/tensorflow/python/data/experimental/ops/matching_files.py
+++ b/tensorflow/python/data/experimental/ops/matching_files.py
@@ -19,9 +19,9 @@
from __future__ import print_function
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
@@ -37,15 +37,5 @@
return ged_ops.experimental_matching_files_dataset(self._patterns)
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.string
-
-
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
diff --git a/tensorflow/python/data/experimental/ops/optimization_options.py b/tensorflow/python/data/experimental/ops/optimization_options.py
index dc9d319..3817e62 100644
--- a/tensorflow/python/data/experimental/ops/optimization_options.py
+++ b/tensorflow/python/data/experimental/ops/optimization_options.py
@@ -26,15 +26,23 @@
class OptimizationOptions(options.OptionsBase):
"""Represents options for dataset optimizations.
- You can apply `OptimizationOptions` to a `dataset` object, as follows:
+ You can set the optimization options of a dataset through the
+ `experimental_optimization` property of `tf.data.Options`; the property is
+ an instance of `tf.data.experimental.OptimizationOptions`.
```python
options = tf.data.Options()
- options.optimization = tf.data.experimental.OptimizationOptions()
- options.optimization.map_and_batch_fusion = True
+ options.experimental_optimization.map_vectorization = True
+ options.apply_default_optimizations = False
dataset = dataset.with_options(options)
```
"""
+ apply_default_optimizations = options.create_option(
+ name="apply_default_optimizations",
+ ty=bool,
+ docstring=
+ "Whether to apply default static optimizations. If False, only static "
+ "optimizations that have been explicitly enabled will be applied.")
filter_fusion = options.create_option(
name="filter_fusion",
@@ -80,4 +88,33 @@
shuffle_and_repeat_fusion = options.create_option(
name="shuffle_and_repeat_fusion",
ty=bool,
- docstring="Whether to fuse shuffle and repeat transformations.")
+ docstring="Whether to fuse shuffle and repeat transformations. If None, "
+ "defaults to True.")
+
+ def _static_optimizations(self):
+ """Produces the list of enabled static optimizations."""
+ result = []
+ optimizations_to_enable = [
+ "filter_fusion",
+ "hoist_random_uniform",
+ "map_and_filter_fusion",
+ "map_fusion",
+ "map_parallelization",
+ "map_vectorization",
+ ]
+ for optimization in optimizations_to_enable:
+ if getattr(self, optimization):
+ result.append(optimization)
+
+ if self.apply_default_optimizations is not False:
+ # The following optimizations are turned on by default, unless the
+ # user explicitly disables them.
+ optimizations_to_disable = [
+ "map_and_batch_fusion",
+ "noop_elimination",
+ "shuffle_and_repeat_fusion",
+ ]
+ for optimization in optimizations_to_disable:
+ if getattr(self, optimization) is not False:
+ result.append(optimization)
+ return result
diff --git a/tensorflow/python/data/experimental/ops/parsing_ops.py b/tensorflow/python/data/experimental/ops/parsing_ops.py
index 44cb0b8..deb20d6 100644
--- a/tensorflow/python/data/experimental/ops/parsing_ops.py
+++ b/tensorflow/python/data/experimental/ops/parsing_ops.py
@@ -18,7 +18,7 @@
from __future__ import print_function
from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
@@ -33,8 +33,8 @@
def __init__(self, input_dataset, features, num_parallel_calls):
super(_ParseExampleDataset, self).__init__(input_dataset)
self._input_dataset = input_dataset
- if not all(types == dtypes.string
- for types in nest.flatten(input_dataset.output_types)):
+ if not input_dataset._element_structure.is_compatible_with( # pylint: disable=protected-access
+ structure.TensorStructure(dtypes.string, [None])):
raise TypeError("Input dataset should be a dataset of vectors of strings")
self._num_parallel_calls = num_parallel_calls
# pylint: disable=protected-access
@@ -67,17 +67,19 @@
for _ in range(len(sparse_keys))
]
- self._output_shapes = dict(
+ output_shapes = dict(
zip(self._dense_keys + self._sparse_keys,
dense_output_shapes + sparse_output_shapes))
- self._output_types = dict(
+ output_types = dict(
zip(self._dense_keys + self._sparse_keys,
self._dense_types + self._sparse_types))
- self._output_classes = dict(
+ output_classes = dict(
zip(self._dense_keys + self._sparse_keys,
[ops.Tensor for _ in range(len(self._dense_defaults))] +
[sparse_tensor.SparseTensor for _ in range(len(self._sparse_keys))
]))
+ self._structure = structure.convert_legacy_structure(
+ output_types, output_shapes, output_classes)
def _as_variant_tensor(self):
return gen_experimental_dataset_ops.experimental_parse_example_dataset(
@@ -91,16 +93,8 @@
**dataset_ops.flat_structure(self))
@property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
-
- @property
- def output_classes(self):
- return self._output_classes
+ def _element_structure(self):
+ return self._structure
# TODO(b/111553342): add arguments names and example names as well.
diff --git a/tensorflow/python/data/experimental/ops/prefetching_ops.py b/tensorflow/python/data/experimental/ops/prefetching_ops.py
index 0894dd5..e3a8622 100644
--- a/tensorflow/python/data/experimental/ops/prefetching_ops.py
+++ b/tensorflow/python/data/experimental/ops/prefetching_ops.py
@@ -19,8 +19,6 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
-from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
from tensorflow.python.eager import function
from tensorflow.python.framework import device as framework_device
from tensorflow.python.framework import dtypes
@@ -73,6 +71,7 @@
def _apply_fn(dataset):
options = dataset_ops.Options()
options.experimental_autotune = False
+ options.experimental_optimization.apply_default_optimizations = False
return _CopyToDeviceDataset(
dataset, target_device=target_device,
source_device=source_device).with_options(options)
@@ -102,13 +101,6 @@
self._source_device_string = source_device
self._source_device = ops.convert_to_tensor(source_device)
- self._flat_output_shapes = nest.flatten(
- sparse.as_dense_shapes(self._input_dataset.output_shapes,
- self._input_dataset.output_classes))
- self._flat_output_types = nest.flatten(
- sparse.as_dense_types(self._input_dataset.output_types,
- self._input_dataset.output_classes))
-
@function.defun()
def _init_func():
"""Creates an iterator for the input dataset.
@@ -119,8 +111,7 @@
# pylint: disable=protected-access
ds_variant = self._input_dataset._as_variant_tensor()
resource = gen_dataset_ops.anonymous_iterator(
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
+ **dataset_ops.flat_structure(self._input_dataset))
with ops.control_dependencies(
[gen_dataset_ops.make_iterator(ds_variant, resource)]):
return gen_dataset_ops.iterator_to_string_handle(resource)
@@ -151,8 +142,7 @@
iterator = iterator_ops.Iterator.from_string_handle(
string_handle, self.output_types, self.output_shapes,
self.output_classes)
- ret = iterator.get_next()
- return nest.flatten(sparse.serialize_sparse_tensors(ret))
+ return self._element_structure._to_tensor_list(iterator.get_next()) # pylint: disable=protected-access
next_func_concrete = _next_func._get_concrete_function_internal() # pylint: disable=protected-access
@@ -162,7 +152,7 @@
target=self._source_device,
args=[string_handle] +
next_func_concrete.captured_inputs,
- Tout=self._flat_output_types,
+ Tout=self._input_dataset._element_structure._flat_types, # pylint: disable=protected-access
f=next_func_concrete)
self._next_func = _remote_next_func._get_concrete_function_internal() # pylint: disable=protected-access
@@ -179,8 +169,7 @@
"""
iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(
string_handle,
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
+ **dataset_ops.flat_structure(self._input_dataset))
with ops.control_dependencies([
resource_variable_ops.destroy_resource_op(
iterator_resource, ignore_lookup_error=True)]):
@@ -192,8 +181,7 @@
def _remote_finalize_func(string_handle):
return functional_ops.remote_call(
target=self._source_device,
- args=[string_handle] +
- finalize_func_concrete.captured_inputs,
+ args=[string_handle] + finalize_func_concrete.captured_inputs,
Tout=[dtypes.int64],
f=finalize_func_concrete)
@@ -229,8 +217,7 @@
init_func=self._init_func,
next_func=self._next_func,
finalize_func=self._finalize_func,
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
+ **dataset_ops.flat_structure(self._input_dataset))
class _MapOnGpuDataset(dataset_ops.UnaryDataset):
@@ -261,16 +248,8 @@
**dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return self._map_func.output_classes
-
- @property
- def output_shapes(self):
- return self._map_func.output_shapes
-
- @property
- def output_types(self):
- return self._map_func.output_types
+ def _element_structure(self):
+ return self._map_func.output_structure
def _transformation_name(self):
return "map_on_gpu()"
diff --git a/tensorflow/python/data/experimental/ops/random_ops.py b/tensorflow/python/data/experimental/ops/random_ops.py
index a4359f3..cbdf367 100644
--- a/tensorflow/python/data/experimental/ops/random_ops.py
+++ b/tensorflow/python/data/experimental/ops/random_ops.py
@@ -21,9 +21,8 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import random_seed
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.util.tf_export import tf_export
@@ -44,16 +43,8 @@
**dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.int64
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.int64, [])
@tf_export(v1=["data.experimental.RandomDataset"])
diff --git a/tensorflow/python/data/experimental/ops/readers.py b/tensorflow/python/data/experimental/ops/readers.py
index a4c260d..c2d82ae 100644
--- a/tensorflow/python/data/experimental/ops/readers.py
+++ b/tensorflow/python/data/experimental/ops/readers.py
@@ -32,10 +32,10 @@
from tensorflow.python.data.ops import readers as core_readers
from tensorflow.python.data.util import convert
from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
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.lib.io import file_io
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.ops import io_ops
@@ -652,11 +652,9 @@
argument_default=[],
argument_dtype=dtypes.int64,
)
- self._output_shapes = tuple(
- tensor_shape.scalar() for _ in range(len(record_defaults)))
- self._output_types = tuple(d.dtype for d in self._record_defaults)
- self._output_classes = tuple(
- ops.Tensor for _ in range(len(record_defaults)))
+ self._structure = structure.NestedStructure(
+ tuple(structure.TensorStructure(d.dtype, [])
+ for d in self._record_defaults))
def _as_variant_tensor(self):
# Constructs graph node for the dataset op.
@@ -665,7 +663,7 @@
record_defaults=self._record_defaults,
buffer_size=self._buffer_size,
header=self._header,
- output_shapes=self._output_shapes,
+ output_shapes=self._structure._flat_shapes, # pylint: disable=protected-access
field_delim=self._field_delim,
use_quote_delim=self._use_quote_delim,
na_value=self._na_value,
@@ -674,16 +672,8 @@
)
@property
- def output_types(self):
- return self._output_types
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_classes(self):
- return self._output_classes
+ def _element_structure(self):
+ return self._structure
@tf_export(v1=["data.experimental.CsvDataset"])
@@ -961,7 +951,9 @@
data_source_name, dtype=dtypes.string, name="data_source_name")
self._query = ops.convert_to_tensor(
query, dtype=dtypes.string, name="query")
- self._output_types = output_types
+ self._structure = structure.NestedStructure(
+ nest.map_structure(
+ lambda dtype: structure.TensorStructure(dtype, []), output_types))
def _as_variant_tensor(self):
return gen_experimental_dataset_ops.experimental_sql_dataset(
@@ -969,17 +961,8 @@
nest.flatten(self.output_types), nest.flatten(self.output_shapes))
@property
- def output_classes(self):
- return nest.map_structure(lambda _: ops.Tensor, self._output_types)
-
- @property
- def output_shapes(self):
- return nest.map_structure(lambda _: tensor_shape.TensorShape([]),
- self._output_types)
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
@tf_export(v1=["data.experimental.SqlDataset"])
diff --git a/tensorflow/python/data/experimental/ops/scan_ops.py b/tensorflow/python/data/experimental/ops/scan_ops.py
index c768373..5c77ad7 100644
--- a/tensorflow/python/data/experimental/ops/scan_ops.py
+++ b/tensorflow/python/data/experimental/ops/scan_ops.py
@@ -21,7 +21,7 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import gen_experimental_dataset_ops
@@ -49,18 +49,7 @@
# Compute initial values for the state classes, shapes and types based on
# the initial state. The shapes may be refined by running `tf_scan_func` one
# or more times below.
- self._state_classes = sparse.get_classes(self._initial_state)
- self._state_shapes = nest.pack_sequence_as(
- self._initial_state,
- [t.get_shape() for t in nest.flatten(self._initial_state)])
- self._state_types = nest.pack_sequence_as(
- self._initial_state,
- [t.dtype for t in nest.flatten(self._initial_state)])
-
- # Will be populated by calling `tf_scan_func`.
- self._output_classes = None
- self._output_shapes = None
- self._output_types = None
+ self._state_structure = structure.Structure.from_value(self._initial_state)
# Iteratively rerun the scan function until reaching a fixed point on
# `self._state_shapes`.
@@ -70,9 +59,8 @@
wrapped_func = dataset_ops.StructuredFunctionWrapper(
scan_func,
self._transformation_name(),
- input_classes=(self._state_classes, input_dataset.output_classes),
- input_shapes=(self._state_shapes, input_dataset.output_shapes),
- input_types=(self._state_types, input_dataset.output_types),
+ input_structure=structure.NestedStructure(
+ (self._state_structure, input_dataset._element_structure)), # pylint: disable=protected-access
add_to_graph=False)
if not (
isinstance(wrapped_func.output_types, collections.Sequence) and
@@ -83,29 +71,35 @@
new_state_classes, self._output_classes = wrapped_func.output_classes
# Extract and validate class information from the returned values.
- for new_state_class, state_class in zip(
+ new_state_classes, output_classes = wrapped_func.output_classes
+ old_state_classes = self._state_structure._to_legacy_output_classes() # pylint: disable=protected-access
+ for new_state_class, old_state_class in zip(
nest.flatten(new_state_classes),
- nest.flatten(self._state_classes)):
- if not issubclass(new_state_class, state_class):
+ nest.flatten(old_state_classes)):
+ if not issubclass(new_state_class, old_state_class):
raise TypeError(
"The element classes for the new state must match the initial "
"state. Expected %s; got %s." %
- (self._state_classes, new_state_classes))
+ (old_state_classes, new_state_classes))
# Extract and validate type information from the returned values.
- new_state_types, self._output_types = wrapped_func.output_types
- for new_state_type, state_type in zip(
- nest.flatten(new_state_types), nest.flatten(self._state_types)):
- if new_state_type != state_type:
+ new_state_types, output_types = wrapped_func.output_types
+ old_state_types = self._state_structure._to_legacy_output_types() # pylint: disable=protected-access
+ for new_state_type, old_state_type in zip(
+ nest.flatten(new_state_types), nest.flatten(old_state_types)):
+ if new_state_type != old_state_type:
raise TypeError(
"The element types for the new state must match the initial "
"state. Expected %s; got %s." %
- (self._state_types, new_state_types))
+ (old_state_types, new_state_types))
# Extract shape information from the returned values.
- new_state_shapes, self._output_shapes = wrapped_func.output_shapes
+ new_state_shapes, output_shapes = wrapped_func.output_shapes
+ old_state_shapes = self._state_structure._to_legacy_output_shapes() # pylint: disable=protected-access
+ self._structure = structure.convert_legacy_structure(
+ output_types, output_shapes, output_classes)
- flat_state_shapes = nest.flatten(self._state_shapes)
+ flat_state_shapes = nest.flatten(old_state_shapes)
flat_new_state_shapes = nest.flatten(new_state_shapes)
weakened_state_shapes = [
original.most_specific_compatible_shape(new)
@@ -122,8 +116,13 @@
break
if need_to_rerun:
- self._state_shapes = nest.pack_sequence_as(self._state_shapes,
- weakened_state_shapes)
+ # TODO(b/110122868): Support a "most specific compatible structure"
+ # method for combining structures, to avoid using legacy structures
+ # in this method.
+ self._state_structure = structure.convert_legacy_structure(
+ old_state_types,
+ nest.pack_sequence_as(old_state_shapes, weakened_state_shapes),
+ old_state_classes)
self._scan_func = wrapped_func
self._scan_func.function.add_to_graph(ops.get_default_graph())
@@ -132,25 +131,19 @@
return [self._scan_func]
def _as_variant_tensor(self):
- input_t = self._input_dataset._as_variant_tensor() # pylint: disable=protected-access
+ # pylint: disable=protected-access
+ input_t = self._input_dataset._as_variant_tensor()
return gen_experimental_dataset_ops.experimental_scan_dataset(
input_t,
- nest.flatten(sparse.serialize_sparse_tensors(self._initial_state)),
+ self._state_structure._to_tensor_list(self._initial_state),
self._scan_func.function.captured_inputs,
f=self._scan_func.function,
+ preserve_cardinality=True,
**dataset_ops.flat_structure(self))
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
def _transformation_name(self):
return "tf.data.experimental.scan()"
diff --git a/tensorflow/python/data/experimental/ops/sleep.py b/tensorflow/python/data/experimental/ops/sleep.py
index 5e9d021..2da8323 100644
--- a/tensorflow/python/data/experimental/ops/sleep.py
+++ b/tensorflow/python/data/experimental/ops/sleep.py
@@ -35,18 +35,6 @@
self._sleep_microseconds,
**dataset_ops.flat_structure(self))
- @property
- def output_shapes(self):
- return self._input_dataset.output_shapes
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
-
- @property
- def output_classes(self):
- return self._input_dataset.output_classes
-
def sleep(sleep_microseconds):
"""Sleeps for `sleep_microseconds` before producing each input element.
diff --git a/tensorflow/python/data/experimental/ops/stats_options.py b/tensorflow/python/data/experimental/ops/stats_options.py
index 6e884aa..94ae67f 100644
--- a/tensorflow/python/data/experimental/ops/stats_options.py
+++ b/tensorflow/python/data/experimental/ops/stats_options.py
@@ -28,27 +28,19 @@
class StatsOptions(options.OptionsBase):
"""Represents options for collecting dataset stats using `StatsAggregator`.
- To apply `StatsOptions` with a `tf.data.Dataset` object, use the following
- pattern:
+ You can set the stats options of a dataset through the `experimental_stats`
+ property of `tf.data.Options`; the property is an instance of
+ `tf.data.experimental.StatsOptions`. For example, to collect latency stats
+ on all dataset edges, use the following pattern:
```python
aggregator = tf.data.experimental.StatsAggregator()
options = tf.data.Options()
- options.experimental_stats = tf.data.experimental.StatsOptions()
options.experimental_stats.aggregator = aggregator
+ options.experimental_stats.latency_all_edges = True
dataset = dataset.with_options(options)
```
-
- Note: a `StatsAggregator` object can be attached either duing construction or
- can be provided later like in above example.
-
- ```python
- aggretator = tf.data.experimental.StatsAggregator()
- # attach aggregator during construction
- options.experimental_stats = tf.data.experimental.StatsOptions(aggregator)
- .....
- ```
"""
aggregator = options.create_option(
@@ -62,18 +54,15 @@
ty=str,
docstring=
"Prefix to prepend all statistics recorded for the input `dataset` with.",
- default="")
+ default_factory=lambda: "")
counter_prefix = options.create_option(
name="counter_prefix",
ty=str,
- docstring=
- "Prefix for the statistics recorded as counter.",
- default="")
+ docstring="Prefix for the statistics recorded as counter.",
+ default_factory=lambda: "")
latency_all_edges = options.create_option(
name="latency_all_edges",
ty=bool,
- docstring=
- "Whether to add latency measurements on all edges.",
- default=True)
+ docstring="Whether to add latency measurements on all edges.")
diff --git a/tensorflow/python/data/experimental/ops/threading_options.py b/tensorflow/python/data/experimental/ops/threading_options.py
index dbf6621..d713b9a 100644
--- a/tensorflow/python/data/experimental/ops/threading_options.py
+++ b/tensorflow/python/data/experimental/ops/threading_options.py
@@ -26,11 +26,12 @@
class ThreadingOptions(options.OptionsBase):
"""Represents options for dataset threading.
- To apply `ThreadingOptions` to a `dataset` object, use the following pattern:
+ You can set the threading options of a dataset through the
+ `experimental_threading` property of `tf.data.Options`; the property is
+ an instance of `tf.data.experimental.ThreadingOptions`.
```python
options = tf.data.Options()
- options.experimental_threading = tf.data.experimental.ThreadingOptions()
options.experimental_threading.private_threadpool_size = 10
dataset = dataset.with_options(options)
```
@@ -46,5 +47,4 @@
name="private_threadpool_size",
ty=int,
docstring=
- "If set, the dataset will use a private threadpool of the given size.",
- default=None)
+ "If set, the dataset will use a private threadpool of the given size.")
diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index 9f7ce99..737ba28 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -97,6 +97,8 @@
"//tensorflow/python/data/util:nest",
"//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
+ "//tensorflow/python:script_ops",
"//tensorflow/python:sparse_tensor",
],
)
@@ -443,6 +445,19 @@
)
tf_py_test(
+ name = "options_test",
+ size = "small",
+ srcs = ["options_test.py"],
+ additional_deps = [
+ ":test_base",
+ "//tensorflow/python/data/ops:dataset_ops",
+ "//tensorflow/python/data/experimental/ops:optimization_options",
+ "//tensorflow/python/data/experimental/ops:threading_options",
+ "//tensorflow/python:client_testlib",
+ ],
+)
+
+tf_py_test(
name = "padded_batch_test",
size = "small",
srcs = ["padded_batch_test.py"],
diff --git a/tensorflow/python/data/kernel_tests/dataset_test.py b/tensorflow/python/data/kernel_tests/dataset_test.py
index 2952c08..820bc8e 100644
--- a/tensorflow/python/data/kernel_tests/dataset_test.py
+++ b/tensorflow/python/data/kernel_tests/dataset_test.py
@@ -207,53 +207,6 @@
self.assertEqual(2, inputs.count(ds2))
self.assertEqual(1, inputs.count(ds3))
- def testOptionsDefault(self):
- ds = dataset_ops.Dataset.range(0)
- self.assertEqual(dataset_ops.Options(), ds.options())
-
- def testOptionsOnce(self):
- options = dataset_ops.Options()
- ds = dataset_ops.Dataset.range(0).with_options(options).cache()
- self.assertEqual(options, ds.options())
-
- def testOptionsTwiceSame(self):
- options = dataset_ops.Options()
- options.experimental_autotune = True
- ds = dataset_ops.Dataset.range(0).with_options(options).with_options(
- options)
- self.assertEqual(options, ds.options())
-
- def testOptionsTwiceDifferent(self):
- options1 = dataset_ops.Options()
- options1.experimental_autotune = True
- options2 = dataset_ops.Options()
- options2.experimental_deterministic = False
- ds = dataset_ops.Dataset.range(0).with_options(options1).with_options(
- options2)
- self.assertTrue(ds.options().experimental_autotune)
- # Explicitly check that flag is False since assertFalse allows None
- self.assertIs(ds.options().experimental_deterministic, False)
-
- def testOptionsTwiceDifferentError(self):
- options1 = dataset_ops.Options()
- options1.experimental_autotune = True
- options2 = dataset_ops.Options()
- options2.experimental_autotune = False
- with self.assertRaisesRegexp(ValueError,
- "Cannot merge incompatible values"):
- dataset_ops.Dataset.range(0).with_options(options1).with_options(options2)
-
- def testOptionsMergeOptionsFromMultipleInputs(self):
- options1 = dataset_ops.Options()
- options1.experimental_autotune = True
- options2 = dataset_ops.Options()
- options2.experimental_deterministic = True
- ds = dataset_ops.Dataset.zip(
- (dataset_ops.Dataset.range(0).with_options(options1),
- dataset_ops.Dataset.range(0).with_options(options2)))
- self.assertTrue(ds.options().experimental_autotune)
- self.assertTrue(ds.options().experimental_deterministic)
-
# TODO(b/119882922): use-after-free bug in eager mode.
# pylint: disable=g-long-lambda
@parameterized.named_parameters(
@@ -313,5 +266,6 @@
round_trip_dataset, [self.evaluate(tf_value_fn())],
requires_initialization=True)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/kernel_tests/from_generator_test.py b/tensorflow/python/data/kernel_tests/from_generator_test.py
index a662553..11919bd 100644
--- a/tensorflow/python/data/kernel_tests/from_generator_test.py
+++ b/tensorflow/python/data/kernel_tests/from_generator_test.py
@@ -21,7 +21,6 @@
import numpy as np
-from tensorflow.python.client import session
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
@@ -32,43 +31,27 @@
from tensorflow.python.platform import test
-class FromGeneratorTest(test_base.DatasetTestBase):
+@test_util.run_all_in_graph_and_eager_modes
+class DatasetConstructorTest(test_base.DatasetTestBase):
def _testFromGenerator(self, generator, elem_sequence, num_repeats,
output_types=None):
if output_types is None:
output_types = dtypes.int64
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(generator, output_types=output_types)
- .repeat(num_repeats)
- .prefetch(5))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- for _ in range(2): # Run twice to test reinitialization.
- sess.run(init_op)
- for _ in range(num_repeats):
- for elem in elem_sequence:
- self.assertAllEqual(elem, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=output_types).repeat(num_repeats).prefetch(5)
+ self.assertDatasetProduces(
+ dataset,
+ elem_sequence * num_repeats,
+ requires_initialization=True,
+ num_test_iterations=2)
def _testFromGeneratorOneShot(self, generator, elem_sequence, num_repeats):
- iterator = dataset_ops.make_one_shot_iterator(
- dataset_ops.Dataset.from_generator(generator, output_types=dtypes.int64)
- .repeat(num_repeats)
- .prefetch(5))
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtypes.int64).repeat(num_repeats).prefetch(5)
+ self.assertDatasetProduces(
+ dataset, elem_sequence * num_repeats, num_test_iterations=2)
- with self.cached_session() as sess:
- for _ in range(num_repeats):
- for elem in elem_sequence:
- self.assertAllEqual(elem, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- @test_util.run_deprecated_v1
def testFromGeneratorUsingFunction(self):
def generator():
for i in range(1, 100):
@@ -79,21 +62,18 @@
self._testFromGeneratorOneShot(generator, elem_sequence, 1)
self._testFromGeneratorOneShot(generator, elem_sequence, 5)
- @test_util.run_deprecated_v1
def testFromGeneratorUsingList(self):
generator = lambda: [[i] * i for i in range(1, 100)]
elem_sequence = list(generator())
self._testFromGenerator(generator, elem_sequence, 1)
self._testFromGenerator(generator, elem_sequence, 5)
- @test_util.run_deprecated_v1
def testFromGeneratorUsingNdarray(self):
generator = lambda: np.arange(100, dtype=np.int64)
elem_sequence = list(generator())
self._testFromGenerator(generator, elem_sequence, 1, output_types=np.int64)
self._testFromGenerator(generator, elem_sequence, 5, output_types=np.int64)
- @test_util.run_deprecated_v1
def testFromGeneratorUsingGeneratorExpression(self):
# NOTE(mrry): Generator *expressions* are not repeatable (or in
# general reusable), because they eagerly evaluate the `for`
@@ -105,7 +85,6 @@
self._testFromGenerator(generator, elem_sequence, 1)
self._testFromGenerator(generator, elem_sequence, 5)
- @test_util.run_deprecated_v1
def testFromMultipleConcurrentGenerators(self):
num_inner_repeats = 5
num_outer_repeats = 100
@@ -128,22 +107,16 @@
output_shapes=([None], [3]))
.repeat(num_inner_repeats).prefetch(5))
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(num_outer_repeats)
- .interleave(interleave_fn, cycle_length=10,
- block_length=len(input_list)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(num_inner_repeats * num_outer_repeats):
- for elem in input_list:
- val0, val1 = sess.run(get_next)
- self.assertAllEqual(elem[0], val0)
- self.assertAllEqual(elem[1], val1)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(num_outer_repeats).interleave(
+ interleave_fn, cycle_length=10, block_length=len(input_list))
+ get_next = self.getNext(dataset)
+ for _ in range(num_inner_repeats * num_outer_repeats):
+ for elem in input_list:
+ val0, val1 = self.evaluate(get_next())
+ self.assertAllEqual(elem[0], val0)
+ self.assertAllEqual(elem[1], val1)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
# TODO(b/67868766): Reenable this when the source of flakiness is discovered.
def _testFromGeneratorsRunningInParallel(self):
@@ -186,22 +159,16 @@
return dataset_ops.Dataset.from_generator(
generator, output_types=dtypes.int64, output_shapes=[]).prefetch(2)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(num_parallel_iterators)
- .interleave(
- interleave_fn, cycle_length=num_parallel_iterators, block_length=1))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.range(num_parallel_iterators).interleave(
+ interleave_fn, cycle_length=num_parallel_iterators, block_length=1)
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- for elem in [0, 1]:
- for _ in range(num_parallel_iterators):
- self.assertAllEqual(elem, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ for elem in [0, 1]:
+ for _ in range(num_parallel_iterators):
+ self.assertAllEqual(elem, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- @test_util.run_deprecated_v1
def testFromGeneratorImplicitConversion(self):
def generator():
yield [1]
@@ -209,45 +176,28 @@
yield [3]
for dtype in [dtypes.int8, dtypes.int32, dtypes.int64]:
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=dtype, output_shapes=[1]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtype, output_shapes=[1])
+ get_next = self.getNext(dataset)
- self.assertEqual(dtype, get_next.dtype)
+ for expected in [[1], [2], [3]]:
+ next_val = self.evaluate(get_next())
+ self.assertEqual(dtype.as_numpy_dtype, next_val.dtype)
+ self.assertAllEqual(expected, next_val)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- with self.cached_session() as sess:
- sess.run(init_op)
- for expected in [[1], [2], [3]]:
- next_val = sess.run(get_next)
- self.assertEqual(dtype.as_numpy_dtype, next_val.dtype)
- self.assertAllEqual(expected, next_val)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- @test_util.run_deprecated_v1
def testFromGeneratorString(self):
def generator():
yield "foo"
yield b"bar"
yield u"baz"
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=dtypes.string, output_shapes=[]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtypes.string, output_shapes=[])
+ self.assertDatasetProduces(
+ dataset, expected_output=[b"foo", b"bar", b"baz"])
- with self.cached_session() as sess:
- sess.run(init_op)
- for expected in [b"foo", b"bar", b"baz"]:
- next_val = sess.run(get_next)
- self.assertAllEqual(expected, next_val)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- @test_util.run_deprecated_v1
def testFromGeneratorTypeError(self):
def generator():
yield np.array([1, 2, 3], dtype=np.int64)
@@ -255,23 +205,19 @@
yield "ERROR"
yield np.array([7, 8, 9], dtype=np.int64)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=dtypes.int64, output_shapes=[3]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtypes.int64, output_shapes=[3])
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual([1, 2, 3], sess.run(get_next))
- self.assertAllEqual([4, 5, 6], sess.run(get_next))
- with self.assertRaisesOpError("The expected type was int64"):
- sess.run(get_next)
- self.assertAllEqual([7, 8, 9], sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ get_next = self.getNext(dataset)
- @test_util.run_deprecated_v1
+ self.assertAllEqual([1, 2, 3], self.evaluate(get_next()))
+ self.assertAllEqual([4, 5, 6], self.evaluate(get_next()))
+ with self.assertRaisesOpError("The expected type was int64"):
+ self.evaluate(get_next())
+ self.assertAllEqual([7, 8, 9], self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+
def testFromGeneratorShapeError(self):
def generator():
yield np.array([1, 2, 3], dtype=np.int64)
@@ -279,23 +225,18 @@
yield np.array([7, 8, 9, 10], dtype=np.int64)
yield np.array([11, 12, 13], dtype=np.int64)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=dtypes.int64, output_shapes=[3]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtypes.int64, output_shapes=[3])
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual([1, 2, 3], sess.run(get_next))
- self.assertAllEqual([4, 5, 6], sess.run(get_next))
- with self.assertRaisesOpError(r"element of shape \(3,\) was expected"):
- sess.run(get_next)
- self.assertAllEqual([11, 12, 13], sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ self.assertAllEqual([1, 2, 3], self.evaluate(get_next()))
+ self.assertAllEqual([4, 5, 6], self.evaluate(get_next()))
+ with self.assertRaisesOpError(r"element of shape \(3,\) was expected"):
+ self.evaluate(get_next())
+ self.assertAllEqual([11, 12, 13], self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- @test_util.run_deprecated_v1
def testFromGeneratorStructureError(self):
def generator():
yield 1, 2
@@ -304,46 +245,31 @@
yield 6, 7, 8
yield 9, 10
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=(dtypes.int64, dtypes.int64)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=(dtypes.int64, dtypes.int64))
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertEqual((1, 2), sess.run(get_next))
- self.assertEqual((3, 4), sess.run(get_next))
- with self.assertRaisesOpError(
- r"The expected structure was \(tf\.int64, tf\.int64\)"):
- sess.run(get_next)
- with self.assertRaisesOpError(
- r"The expected structure was \(tf\.int64, tf\.int64\)"):
- sess.run(get_next)
- self.assertEqual((9, 10), sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ self.assertEqual((1, 2), self.evaluate(get_next()))
+ self.assertEqual((3, 4), self.evaluate(get_next()))
+ with self.assertRaisesOpError(
+ r"The expected structure was \(tf\.int64, tf\.int64\)"):
+ self.evaluate(get_next())
+ with self.assertRaisesOpError(
+ r"The expected structure was \(tf\.int64, tf\.int64\)"):
+ self.evaluate(get_next())
+ self.assertEqual((9, 10), self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- @test_util.run_deprecated_v1
def testFromGeneratorHeterogeneous(self):
def generator():
yield 1
yield [2, 3]
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=dtypes.int64))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtypes.int64)
+ self.assertDatasetProduces(dataset, expected_output=[1, [2, 3]])
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual(1, sess.run(get_next))
- self.assertAllEqual([2, 3], sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- @test_util.run_deprecated_v1
def testFromGeneratorStopShort(self):
def generator():
@@ -351,18 +277,12 @@
yield 1
yield 2
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- generator, output_types=dtypes.int64))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ generator, output_types=dtypes.int64)
+ get_next = self.getNext(dataset)
+ self.assertAllEqual(0, self.evaluate(get_next()))
+ self.assertAllEqual(1, self.evaluate(get_next()))
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual(0, sess.run(get_next))
- self.assertAllEqual(1, sess.run(get_next))
-
- @test_util.run_deprecated_v1
def testFromGeneratorDestructorCalled(self):
# Use an `Event` to signal that the generator has been deleted.
event = threading.Event()
@@ -381,23 +301,18 @@
def __del__(self):
event.set()
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_generator(
- GeneratorWrapper, output_types=dtypes.int64).take(2))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_generator(
+ GeneratorWrapper, output_types=dtypes.int64).take(2)
+ get_next = self.getNext(dataset)
- with session.Session() as sess:
- sess.run(init_op)
- self.assertAllEqual(42, sess.run(get_next))
- self.assertAllEqual(42, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
- # Test that `GeneratorWrapper` object is destroyed when the
- # iterator terminates (and the generator iterator is deleted).
- self.assertTrue(event.is_set())
+ self.assertAllEqual(42, self.evaluate(get_next()))
+ self.assertAllEqual(42, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+ # Test that `GeneratorWrapper` object is destroyed when the
+ # iterator terminates (and the generator iterator is deleted).
+ self.assertTrue(event.is_set())
- @test_util.run_deprecated_v1
def testFromGeneratorWithArgs(self):
def flat_map_fn(elem):
@@ -410,20 +325,10 @@
generator_with_arg, output_types=dtypes.int64, output_shapes=(),
args=(elem,))
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(5).flat_map(flat_map_fn))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.range(5).flat_map(flat_map_fn)
+ self.assertDatasetProduces(
+ dataset, expected_output=[1, 2, 2, 3, 3, 3, 4, 4, 4, 4])
- with self.cached_session() as sess:
- sess.run(init_op)
- expected = [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]
- for x in expected:
- self.assertEqual(x, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- @test_util.run_deprecated_v1
def testFromGeneratorWithTwoArgs(self):
def flat_map_fn(elem, message):
@@ -436,26 +341,17 @@
generator_with_arg, output_types=(dtypes.int64, dtypes.string),
output_shapes=((), ()), args=(elem, message))
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.zip(
- (dataset_ops.Dataset.range(5),
- dataset_ops.Dataset.from_tensors("Hi!").repeat(None)))
- .flat_map(flat_map_fn))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.zip(
+ (dataset_ops.Dataset.range(5),
+ dataset_ops.Dataset.from_tensors("Hi!").repeat(None)
+ )).flat_map(flat_map_fn)
- with self.cached_session() as sess:
- sess.run(init_op)
- expected = [(0, b"Hi!"),
- (0, b"Hi!"), (1, b"Hi!"),
- (0, b"Hi!"), (1, b"Hi!"), (2, b"Hi!"),
- (0, b"Hi!"), (1, b"Hi!"), (2, b"Hi!"), (3, b"Hi!")]
- for x in expected:
- self.assertEqual(x, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ self.assertDatasetProduces(
+ dataset,
+ expected_output=[(0, b"Hi!"), (0, b"Hi!"), (1, b"Hi!"), (0, b"Hi!"),
+ (1, b"Hi!"), (2, b"Hi!"), (0, b"Hi!"), (1, b"Hi!"),
+ (2, b"Hi!"), (3, b"Hi!")])
- @test_util.run_deprecated_v1
def testGeneratorDatasetFinalizeFunctionCalled(self):
# NOTE(mrry): This test tests the internal `_GeneratorDataset`,
# which affords more control over what the finalize function can do than
@@ -472,19 +368,15 @@
stateful=True)
dummy = constant_op.constant(37)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops._GeneratorDataset(
- dummy, lambda x: x, lambda x: x, finalize_fn).take(2))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops._GeneratorDataset(dummy, lambda x: x, lambda x: x,
+ finalize_fn).take(2)
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual(37, sess.run(get_next))
- self.assertAllEqual(37, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
- self.assertTrue(event.is_set())
+ self.assertAllEqual(37, self.evaluate(get_next()))
+ self.assertAllEqual(37, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+ self.assertTrue(event.is_set())
if __name__ == "__main__":
diff --git a/tensorflow/python/data/kernel_tests/interleave_test.py b/tensorflow/python/data/kernel_tests/interleave_test.py
index c3450e6..05a211a 100644
--- a/tensorflow/python/data/kernel_tests/interleave_test.py
+++ b/tensorflow/python/data/kernel_tests/interleave_test.py
@@ -264,6 +264,7 @@
("8", np.int64([4, 0, 6]), 2, 3, 1),
("9", np.int64([4, 0, 6]), 2, 3, 2),
)
+ @test_util.run_v1_only("b/120545219")
def testSkipEagerSloppyInterleaveInOrder(self, input_values, cycle_length,
block_length, num_parallel_calls):
get_next, coordination_events = _make_coordinated_sloppy_dataset(
@@ -286,6 +287,7 @@
("3", np.int64([4, 5, 6]), 3, 2, 3),
("4", np.int64([4, 0, 6]), 2, 3, 2),
)
+ @test_util.run_v1_only("b/120545219")
def testSkipEagerSloppyInterleaveOutOfOrder(self, input_values, cycle_length,
block_length, num_parallel_calls):
get_next, coordination_events = _make_coordinated_sloppy_dataset(
diff --git a/tensorflow/python/data/kernel_tests/iterator_cluster_test.py b/tensorflow/python/data/kernel_tests/iterator_cluster_test.py
index 728bed2..2008823 100644
--- a/tensorflow/python/data/kernel_tests/iterator_cluster_test.py
+++ b/tensorflow/python/data/kernel_tests/iterator_cluster_test.py
@@ -39,6 +39,7 @@
class IteratorClusterTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testRemoteIteratorWithoutRemoteCallFail(self):
worker_config = config_pb2.ConfigProto()
worker_config.device_count["CPU"] = 2
@@ -92,6 +93,7 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(remote_op, feed_dict={target_placeholder: device1})
+ @test_util.run_v1_only("b/120545219")
def testRemoteIteratorUsingRemoteCallOp(self):
worker_config = config_pb2.ConfigProto()
worker_config.device_count["CPU"] = 2
@@ -102,6 +104,7 @@
"/job:worker/replica:0/task:0/cpu:1",
worker[0].target)
+ @test_util.run_v1_only("b/120545219")
def testRemoteIteratorUsingRemoteCallOpCrossProcess(self):
workers, _ = test_util.create_local_cluster(2, 1)
@@ -109,6 +112,7 @@
"/job:worker/replica:0/task:1/cpu:0",
workers[0].target)
+ @test_util.run_v1_only("b/120545219")
def testCaptureHashTableInSharedIterator(self):
worker, _ = test_util.create_local_cluster(1, 1)
@@ -143,6 +147,7 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
+ @test_util.run_v1_only("b/120545219")
def testImplicitDisposeParallelMapDataset(self):
# Tests whether a parallel map dataset will be cleaned up correctly when
# the pipeline does not run it until exhaustion.
diff --git a/tensorflow/python/data/kernel_tests/list_files_test.py b/tensorflow/python/data/kernel_tests/list_files_test.py
index 789f1ab..a70c4b0 100644
--- a/tensorflow/python/data/kernel_tests/list_files_test.py
+++ b/tensorflow/python/data/kernel_tests/list_files_test.py
@@ -82,27 +82,29 @@
dataset = dataset_ops.Dataset.list_files(
path.join(self.tmp_dir, '*'), shuffle=True, seed=37)
- full_filenames = [compat.as_bytes(path.join(self.tmp_dir, filename))
- for filename in filenames]
+ expected_filenames = [
+ compat.as_bytes(path.join(self.tmp_dir, filename))
+ for filename in filenames
+ ]
- all_produced_filenames = []
+ all_actual_filenames = []
for _ in range(3):
- produced_filenames = []
+ actual_filenames = []
next_element = self.getNext(dataset, requires_initialization=True)
try:
while True:
- produced_filenames.append(self.evaluate(next_element()))
+ actual_filenames.append(self.evaluate(next_element()))
except errors.OutOfRangeError:
pass
- all_produced_filenames.append(produced_filenames)
+ all_actual_filenames.append(actual_filenames)
# Each run should produce the same set of filenames, which may be
- # different from the order of `full_filenames`.
- self.assertItemsEqual(full_filenames, all_produced_filenames[0])
+ # different from the order of `expected_filenames`.
+ self.assertItemsEqual(expected_filenames, all_actual_filenames[0])
# However, the different runs should produce filenames in the same order
# as each other.
- self.assertEqual(all_produced_filenames[0], all_produced_filenames[1])
- self.assertEqual(all_produced_filenames[0], all_produced_filenames[2])
+ self.assertEqual(all_actual_filenames[0], all_actual_filenames[1])
+ self.assertEqual(all_actual_filenames[0], all_actual_filenames[2])
# TODO(b/117581999): eager mode assertion fail wrapped, debug.
def tesSkipEagerEmptyDirectoryInitializer(self):
@@ -169,16 +171,17 @@
path.join(self.tmp_dir, '*'), shuffle=False).repeat(2)
next_element = self.getNext(dataset)
- full_filenames = []
- produced_filenames = []
+ expected_filenames = []
+ actual_filenames = []
for filename in filenames * 2:
- full_filenames.append(compat.as_bytes(path.join(self.tmp_dir, filename)))
- produced_filenames.append(compat.as_bytes(self.evaluate(next_element())))
+ expected_filenames.append(
+ compat.as_bytes(path.join(self.tmp_dir, filename)))
+ actual_filenames.append(compat.as_bytes(self.evaluate(next_element())))
with self.assertRaises(errors.OutOfRangeError):
self.evaluate(next_element())
- self.assertItemsEqual(full_filenames, produced_filenames)
- self.assertEqual(produced_filenames[:len(filenames)],
- produced_filenames[len(filenames):])
+ self.assertItemsEqual(expected_filenames, actual_filenames)
+ self.assertEqual(actual_filenames[:len(filenames)],
+ actual_filenames[len(filenames):])
def testMultiplePatternsAsList(self):
filenames = ['a.txt', 'b.py', 'c.py', 'd.pyc']
diff --git a/tensorflow/python/data/kernel_tests/map_test.py b/tensorflow/python/data/kernel_tests/map_test.py
index fdce791..67ef98f 100644
--- a/tensorflow/python/data/kernel_tests/map_test.py
+++ b/tensorflow/python/data/kernel_tests/map_test.py
@@ -28,12 +28,14 @@
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
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 sparse_tensor
from tensorflow.python.framework import tensor_util
+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 data_flow_ops
@@ -45,6 +47,7 @@
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
from tensorflow.python.platform import test
@@ -82,13 +85,19 @@
return next_element, coordination_events
+@test_util.run_all_in_graph_and_eager_modes
class MapDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
def _buildMapDataset(self, components, count):
+
def _map_fn(x, y, z):
return math_ops.square(x), math_ops.square(y), math_ops.square(z)
- return (dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn)
- .repeat(count))
+
+ dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+ _map_fn).repeat(count)
+ self.assertEqual([c.shape[1:] for c in components],
+ [shape for shape in dataset.output_shapes])
+ return dataset
def testMapDataset(self):
"""Test an dataset that maps a TF function across its input elements."""
@@ -97,34 +106,32 @@
components = (np.arange(7),
np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
np.array(37.0) * np.arange(7))
- count = array_ops.placeholder(dtypes.int64, shape=[])
- dataset = self._buildMapDataset(components, count)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ # Test single-threaded access to the iterator.
+ get_next = self.getNext(self._buildMapDataset(components, 14))
+ for _ in range(14):
+ for i in range(7):
+ result = self.evaluate(get_next())
+ for component, result_component in zip(components, result):
+ self.assertAllEqual(component[i]**2, result_component)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- self.assertEqual([c.shape[1:] for c in components],
- [t.shape for t in get_next])
-
+ # TODO(b/117581999): add eager coverage, different threads run in graph
+ # context.
+ @test_util.run_v1_only("b/120545219")
+ def testSkipEagerMapDatasetMultithreaded(self):
+ # Test multi-threaded access to the same iterator.
+ components = (np.arange(7),
+ np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
+ np.array(37.0) * np.arange(7))
+ get_next = self.getNext(self._buildMapDataset(components, 18))
+ results = []
with self.cached_session() as sess:
- # Test single-threaded access to the iterator.
- sess.run(init_op, feed_dict={count: 14})
- for _ in range(14):
- for i in range(7):
- result = sess.run(get_next)
- for component, result_component in zip(components, result):
- self.assertAllEqual(component[i]**2, result_component)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
-
- # Test multi-threaded access to the same iterator.
- sess.run(init_op, feed_dict={count: 18})
- results = []
def iterator_thread():
while True:
try:
- results.append(sess.run(get_next))
+ results.append(sess.run(get_next()))
except errors.OutOfRangeError:
return
threads = [self.checkedThread(target=iterator_thread) for _ in range(8)]
@@ -146,59 +153,66 @@
def _buildParallelMapDataset(self, components, count, num_parallel_calls,
output_buffer_size):
+
def _map_fn(x, y, z):
return math_ops.square(x), math_ops.square(y), math_ops.square(z)
- return (dataset_ops.Dataset.from_tensor_slices(components)
- .map(_map_fn, num_parallel_calls=num_parallel_calls)
- .prefetch(output_buffer_size)
- .repeat(count))
+
+ dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+ _map_fn, num_parallel_calls=num_parallel_calls).prefetch(
+ output_buffer_size).repeat(count)
+
+ self.assertEqual([c.shape[1:] for c in components],
+ [shape for shape in dataset.output_shapes])
+ return dataset
def testParallelMapDataset(self):
"""Test an dataset that maps a TF function across its input elements."""
+
# The pipeline is TensorSliceDataset -> ParallelMapDataset(square_3) ->
# RepeatDataset(count).
- components = (np.arange(7),
- np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
- np.array(37.0) * np.arange(7))
- count = array_ops.placeholder(dtypes.int64, shape=[])
- num_parallel_calls = array_ops.placeholder(dtypes.int32, shape=[])
- output_buffer_size = array_ops.placeholder(dtypes.int64, shape=[])
+ def do_test(num_parallel_calls, output_buffer_size):
- dataset = self._buildParallelMapDataset(
- components, count, num_parallel_calls, output_buffer_size)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ components = (np.arange(7),
+ np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
+ np.array(37.0) * np.arange(7))
+ # Test single-threaded access to the iterator.
+ get_next = self.getNext(
+ self._buildParallelMapDataset(components, 14, num_parallel_calls,
+ output_buffer_size))
+ for _ in range(14):
+ for i in range(7):
+ result = self.evaluate(get_next())
+ for component, result_component in zip(components, result):
+ self.assertAllEqual(component[i]**2, result_component)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- self.assertEqual([c.shape[1:] for c in components],
- [t.shape for t in get_next])
+ for num_parallel_calls_val, output_buffer_size_val in [(1, 1), (1, 2), (2,
+ 2),
+ (2, 4), (8, 8),
+ (8, 16)]:
+ do_test(num_parallel_calls_val, output_buffer_size_val)
- with self.cached_session() as sess:
+ # TODO(b/117581999): add eager coverage, different threads run in graph
+ # context.
+ @test_util.run_v1_only("b/120545219")
+ def testSkipEagerParallelMapDatasetMultithreaded(self):
- def do_test(num_parallel_calls_val, output_buffer_size_val):
- # Test single-threaded access to the iterator.
- sess.run(init_op, feed_dict={
- count: 14,
- num_parallel_calls: num_parallel_calls_val,
- output_buffer_size: output_buffer_size_val})
- for _ in range(14):
- for i in range(7):
- result = sess.run(get_next)
- for component, result_component in zip(components, result):
- self.assertAllEqual(component[i]**2, result_component)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ def do_test(num_parallel_calls, output_buffer_size):
+ # Test multi-threaded access to the same iterator.
+ components = (np.arange(7),
+ np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
+ np.array(37.0) * np.arange(7))
+ get_next = self.getNext(
+ self._buildParallelMapDataset(components, 18, num_parallel_calls,
+ output_buffer_size))
+ results = []
+ with self.cached_session() as sess:
- # Test multi-threaded access to the same iterator.
- sess.run(init_op, feed_dict={
- count: 18,
- num_parallel_calls: num_parallel_calls_val,
- output_buffer_size: output_buffer_size_val})
- results = []
def iterator_thread():
while True:
try:
- results.append(sess.run(get_next))
+ results.append(sess.run(get_next()))
except errors.OutOfRangeError:
return
threads = [self.checkedThread(target=iterator_thread)
@@ -235,14 +249,10 @@
dataset = self._buildParallelMapDataset(components, 1000, 100, 100)
# NOTE(mrry): Also test that the prefetching thread is cancelled correctly.
dataset = dataset.prefetch(100)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(3):
- sess.run(get_next)
+ for _ in range(3):
+ self.evaluate(get_next())
def testParallelMapUnspecifiedOutputSize(self):
components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
@@ -250,14 +260,10 @@
dataset = (dataset_ops.Dataset.from_tensor_slices(components)
.map(lambda x: array_ops.check_numerics(x, "message"),
num_parallel_calls=2))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(3):
- sess.run(get_next)
+ for _ in range(3):
+ self.evaluate(get_next())
def testParallelMapError(self):
components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
@@ -265,20 +271,16 @@
dataset = (dataset_ops.Dataset.from_tensor_slices(components)
.map(lambda x: array_ops.check_numerics(x, "message"),
num_parallel_calls=2))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(3):
- sess.run(get_next)
- # The 4th element is NaN, so `array_ops.check_numerics()` should fail.
- with self.assertRaises(errors.InvalidArgumentError):
- sess.run(get_next)
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ for _ in range(3):
+ self.evaluate(get_next())
+ # The 4th element is NaN, so `array_ops.check_numerics()` should fail.
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(get_next())
+ self.evaluate(get_next())
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testPrefetchError(self):
components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
@@ -286,20 +288,17 @@
dataset = (dataset_ops.Dataset.from_tensor_slices(components)
.map(lambda x: array_ops.check_numerics(x, "message"))
.prefetch(2))
- iterator = dataset_ops.make_initializable_iterator(dataset)
- init_op = iterator.initializer
- get_next = iterator.get_next()
- with self.cached_session() as sess:
- sess.run(init_op)
- for _ in range(3):
- sess.run(get_next)
- # The 4th element is NaN, so `array_ops.check_numerics()` should fail.
- with self.assertRaises(errors.InvalidArgumentError):
- sess.run(get_next)
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ get_next = self.getNext(dataset)
+
+ for _ in range(3):
+ self.evaluate(get_next())
+ # The 4th element is NaN, so `array_ops.check_numerics()` should fail.
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(get_next())
+ self.evaluate(get_next())
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testCaptureIterator(self):
@@ -312,23 +311,22 @@
return dataset_ops.Dataset.range(10).map(_map_fn)
def _build_graph():
- captured_iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10))
+ if context.executing_eagerly():
+ captured_iterator = iter(dataset_ops.Dataset.range(10))
+ else:
+ captured_iterator = dataset_ops.Dataset.range(
+ 10).make_initializable_iterator()
ds = _build_ds(captured_iterator)
- iterator = ds.make_initializable_iterator()
- init_op = iterator.initializer
- get_next = iterator.get_next()
- return captured_iterator.initializer, init_op, get_next
+ return captured_iterator, ds
- with ops.Graph().as_default() as g:
- captured_init_op, init_op, get_next = _build_graph()
- with self.session(graph=g) as sess:
- sess.run(captured_init_op)
- sess.run(init_op)
- for i in range(10):
- self.assertEqual(i * i, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ captured_iter, ds = _build_graph()
+ if not context.executing_eagerly():
+ self.evaluate(captured_iter.initializer)
+ get_next = self.getNext(ds, requires_initialization=True)
+ for i in range(10):
+ self.assertEqual(i * i, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testCaptureHashTable(self):
# NOTE(mrry): We must use the V2 variants of `HashTable`
@@ -343,41 +341,37 @@
input_sentences = dataset_ops.Dataset.from_tensor_slices(
["brain brain tank salad surgery", "surgery brain"])
- iterator = dataset_ops.make_initializable_iterator(
- input_sentences
- .map(lambda x: string_ops.string_split([x]).values).map(table.lookup))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = input_sentences.map(lambda x: string_ops.string_split([x]).values
+ ).map(table.lookup)
- with self.cached_session() as sess:
- sess.run(table.initializer)
- sess.run(init_op)
- sess.run(get_next)
- sess.run(get_next)
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ get_next = self.getNext(dataset, requires_initialization=True)
+
+ self.evaluate(table.initializer)
+ self.evaluate(get_next())
+ self.evaluate(get_next())
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testCaptureQueue(self):
elements = np.random.randint(100, size=[200])
queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
enqueue_op = queue.enqueue_many(elements)
close_op = queue.close()
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(0).repeat(-1)
- .map(lambda _: queue.dequeue()))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors(0).repeat(
+ -1).map(lambda _: queue.dequeue())
- with self.cached_session() as sess:
- sess.run(enqueue_op)
- sess.run(close_op)
- sess.run(init_op)
- for element in elements:
- self.assertEqual(element, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ get_next = self.getNext(dataset, requires_initialization=True)
+ self.evaluate(enqueue_op)
+ self.evaluate(close_op)
- def testCaptureSameResourceMultipleTimes(self):
+ for element in elements:
+ self.assertEqual(element, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+
+ # TODO(b/117581999): Possible deadlock in eager mode, debug.
+ @test_util.run_v1_only("b/120545219")
+ def testSkipEagerCaptureSameResourceMultipleTimes(self):
elements = np.random.randint(100, size=[200])
queue = data_flow_ops.FIFOQueue(
200, dtypes.int64, shapes=[], shared_name="shared_queue")
@@ -387,101 +381,84 @@
enqueue_op = queue.enqueue_many(elements)
close_op = queue.close()
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(0).repeat(-1)
- .map(lambda _: (queue.dequeue(), queue_2.dequeue())))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors(0).repeat(
+ -1).map(lambda _: (queue.dequeue(), queue_2.dequeue()))
- with self.cached_session() as sess:
- sess.run(enqueue_op)
- sess.run(close_op)
- sess.run(init_op)
- for i in range(100):
- self.assertEqual(sorted([elements[i * 2], elements[i * 2 + 1]]),
- sorted(sess.run(get_next)))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ self.evaluate(enqueue_op)
+ self.evaluate(close_op)
+ get_next = self.getNext(dataset, requires_initialization=True)
+ for i in range(100):
+ self.assertCountEqual([elements[i * 2], elements[i * 2 + 1]],
+ self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testCaptureVariable(self):
counter_var = variable_scope.get_variable(
"counter", (), dtypes.int32, use_resource=True)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(0).repeat(10)
- .map(lambda _: counter_var.assign_add(1)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors(0).repeat(
+ 10).map(lambda _: counter_var.assign_add(1))
+ get_next = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- sess.run(counter_var.initializer)
- sess.run(init_op)
- for i in range(10):
- self.assertEqual(i, sess.run(counter_var))
- self.assertEqual(i + 1, sess.run(get_next))
- self.assertEqual(10, sess.run(counter_var))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
- self.assertEqual(10, sess.run(counter_var))
+ self.evaluate(counter_var.initializer)
- def testCaptureUninitializedVariableError(self):
+ for i in range(10):
+ self.assertEqual(i, self.evaluate(counter_var))
+ self.assertEqual(i + 1, self.evaluate(get_next()))
+ self.assertEqual(10, self.evaluate(counter_var))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+ self.assertEqual(10, self.evaluate(counter_var))
+
+ # TODO(b/117581999): error not captured for eager mode, debug.
+ @test_util.run_v1_only("b/120545219")
+ def testSkipEagerCaptureUninitializedVariableError(self):
counter_var = variable_scope.get_variable(
"counter", (), dtypes.int32, use_resource=True)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(0).repeat(10)
- .map(lambda _: counter_var.assign_add(1)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors(0).repeat(
+ 10).map(lambda _: counter_var.assign_add(1))
- with self.cached_session() as sess:
- sess.run(init_op)
- with self.assertRaises(errors.NotFoundError):
- sess.run(get_next)
+ get_next = self.getNext(dataset, requires_initialization=True)
+
+ with self.assertRaises(errors.NotFoundError):
+ self.evaluate(get_next())
def testSeededStatefulOperatorIsProperlyStateful(self):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(0).repeat(10)
- .map(lambda _: random_ops.random_uniform((), seed=11)).batch(2))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors(0).repeat(
+ 10).map(lambda _: random_ops.random_uniform((), seed=11)).batch(2)
- with self.cached_session() as sess:
- sess.run(init_op)
- random_values = []
- with self.assertRaises(errors.OutOfRangeError):
- while True:
- random_values.extend(sess.run(get_next))
- self.assertEqual(10, len(random_values))
- self.assertGreater(np.abs(np.diff(random_values)).max(), 1e-6)
- sess.run(init_op)
- random_values_2 = []
- with self.assertRaises(errors.OutOfRangeError):
- while True:
- random_values_2.extend(sess.run(get_next))
+ get_next = self.getNext(dataset, requires_initialization=True)
+ random_values = []
+ with self.assertRaises(errors.OutOfRangeError):
+ while True:
+ random_values.extend(self.evaluate(get_next()))
+ self.assertLen(random_values, 10)
+ self.assertGreater(np.abs(np.diff(random_values)).max(), 1e-6)
- # Randomness is repeatable given same seed
- self.assertAllClose(random_values, random_values_2)
+ get_next = self.getNext(dataset, requires_initialization=True)
+ random_values_2 = []
+ with self.assertRaises(errors.OutOfRangeError):
+ while True:
+ random_values_2.extend(self.evaluate(get_next()))
+
+ # Randomness is repeatable given same seed
+ self.assertAllClose(random_values, random_values_2)
def testStatefulMapKeepsStateAcrossIterators(self):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(0).repeat(10)
- .map(lambda _: random_ops.random_uniform((), seed=11))
- .repeat(1000)
- .batch(10))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.from_tensors(0).repeat(10).map(
+ lambda _: random_ops.random_uniform((), seed=11)).repeat(1000).batch(10)
- with self.cached_session() as sess:
- sess.run(init_op)
- random_values = sess.run(get_next)
+ get_next = self.getNext(dataset)
+ random_values = self.evaluate(get_next())
- # Assert that one of the next 99 batches yielded by the iterator is
- # different from the first.
- i = 0
- while i < 99:
- if np.any(random_values != sess.run(get_next)):
- break
- i += 1
- self.assertLess(i, 99)
+ # Assert that one of the next 99 batches yielded by the iterator is
+ # different from the first.
+ i = 0
+ while i < 99:
+ if np.any(random_values != self.evaluate(get_next())):
+ break
+ i += 1
+ self.assertLess(i, 99)
def testStatefulOperationInShortCircuit(self):
counter_var = variable_scope.get_variable(
@@ -491,36 +468,25 @@
counter_var.assign_add(1)
return x
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10).map(increment_fn))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.range(10).map(increment_fn)
- with self.cached_session() as sess:
- sess.run(counter_var.initializer)
- sess.run(init_op)
- for i in range(10):
- self.assertEqual(i, sess.run(counter_var))
- self.assertEqual(i, sess.run(get_next))
- self.assertEqual(10, sess.run(counter_var))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
- self.assertEqual(10, sess.run(counter_var))
+ get_next = self.getNext(dataset, requires_initialization=True)
+
+ self.evaluate(counter_var.initializer)
+ for i in range(10):
+ self.assertEqual(i, self.evaluate(counter_var))
+ self.assertEqual(i, self.evaluate(get_next()))
+ self.assertEqual(10, self.evaluate(counter_var))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+ self.assertEqual(10, self.evaluate(counter_var))
def testMapDict(self):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10)
- .map(lambda x: {"foo": x * 2, "bar": x ** 2})
- .map(lambda d: d["foo"] + d["bar"]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(10):
- self.assertEqual(i * 2 + i**2, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(10).map(
+ lambda x: {"foo": x * 2, "bar": x**2}).map(
+ lambda d: d["foo"] + d["bar"])
+ self.assertDatasetProduces(
+ dataset, expected_output=[i * 2 + i**2 for i in range(10)])
def testMapNamedtuple(self, count=10):
# construct dataset of tuples
@@ -543,33 +509,23 @@
dataset_tuple = dataset_tuple.map(preprocess_tuple)
dataset_namedtuple = dataset_namedtuple.map(preprocess_namedtuple)
- next_tuple = dataset_ops.make_one_shot_iterator(dataset_tuple).get_next()
- next_namedtuple = dataset_ops.make_one_shot_iterator(
- dataset_namedtuple).get_next()
+ next_tuple = self.getNext(dataset_tuple)
+ next_namedtuple = self.getNext(dataset_namedtuple)
# make sure both datasets contain the same data
- with self.cached_session() as sess:
- for i in range(count):
- tuple_, namedtuple_ = sess.run([next_tuple, next_namedtuple])
- self.assertEqual(tuple_, namedtuple_)
- self.assertEqual(tuple_, (i, -2 * i))
+ for i in range(count):
+ tuple_, namedtuple_ = self.evaluate([next_tuple(), next_namedtuple()])
+ self.assertEqual(tuple_, namedtuple_)
+ self.assertEqual(tuple_, (i, -2 * i))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(next_namedtuple)
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(next_namedtuple())
def testUseStepContainerInMap(self):
row = np.arange(6)
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(row)
- .map(lambda elems: functional_ops.map_fn(lambda x: x * x, elems)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual(row**2, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.from_tensors(
+ row).map(lambda elems: functional_ops.map_fn(lambda x: x * x, elems))
+ self.assertDatasetProduces(dataset, expected_output=[row**2])
def testCaseAndCondInMap(self):
@@ -597,24 +553,19 @@
pred_fn_pairs, default=multiply, exclusive=True)
def build_dataset(row, num):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensor_slices(row).map(
- lambda x: control_map_fn(x, num)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
- return init_op, get_next
+ dataset = dataset_ops.Dataset.from_tensor_slices(
+ row).map(lambda x: control_map_fn(x, num))
+ return self.getNext(dataset)
- with self.cached_session() as sess:
- row = np.arange(6)
- for num in [2, 3, 4]:
- init_op, get_next = build_dataset(row, num)
- sess.run(init_op)
- for i in range(6):
- self.assertEqual(
- (i // 2 if i % 2 else i * 2) if (num == 2 or num == 3) else i * 2,
- sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ row = np.arange(6)
+ for num in [2, 3, 4]:
+ get_next = build_dataset(row, num)
+ for i in range(6):
+ self.assertEqual(
+ (i // 2 if i % 2 else i * 2) if (num == 2 or num == 3) else i * 2,
+ self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testCaseInWhileInMap(self):
@@ -636,24 +587,19 @@
def build_dataset(row, num):
# pylint: disable=g-long-lambda
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(row).map(
- lambda elems: functional_ops.map_fn(
- lambda x: control_map_fn(x, num), elems)))
- init_op = iterator.initializer
- get_next = iterator.get_next()
- return init_op, get_next
+ dataset = dataset_ops.Dataset.from_tensors(
+ row).map(lambda elems: functional_ops.map_fn(
+ lambda x: control_map_fn(x, num), elems))
+ return self.getNext(dataset)
- with self.cached_session() as sess:
- row = np.arange(6)
- for num in [2, 3, 4]:
- init_op, get_next = build_dataset(row, num)
- sess.run(init_op)
- self.assertAllEqual(
- [x // 2 if (num == 2 or num == 3) else x * 2 for x in row],
- sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ row = np.arange(6)
+ for num in [2, 3, 4]:
+ get_next = build_dataset(row, num)
+ self.assertAllEqual(
+ [x // 2 if (num == 2 or num == 3) else x * 2 for x in row],
+ self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testCaseAndCondInWhileInMap(self):
@@ -683,21 +629,17 @@
row = np.arange(6)
num = 2
# pylint: disable=g-long-lambda
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.from_tensors(row).map(
- lambda elems: functional_ops.map_fn(
- lambda x: control_map_fn(x, num), elems)))
+ dataset = dataset_ops.Dataset.from_tensors(
+ row).map(lambda elems: functional_ops.map_fn(
+ lambda x: control_map_fn(x, num), elems))
# pylint: enable=g-long-lambda
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- sess.run(init_op)
- self.assertAllEqual([(x // 2 if x % 2 else x * 2) if
- (num == 2 or num == 3) else x * 2 for x in row],
- sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ self.assertAllEqual([(x // 2 if x % 2 else x * 2) if
+ (num == 2 or num == 3) else x * 2 for x in row],
+ self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testPrefetch(self):
# We will use this event to test that `_map_py_func()` has been
@@ -715,58 +657,54 @@
def _map_fn(x):
return script_ops.py_func(_map_py_func, [x], x.dtype)
- buffer_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[])
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(100)
- .map(_map_fn)
- .prefetch(buffer_size_placeholder))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ def do_test(buffer_size):
+ dataset = dataset_ops.Dataset.range(100).map(_map_fn).prefetch(
+ buffer_size)
- with self.cached_session() as sess:
+ get_next = self.getNext(dataset)
# Simple test that prefetch yields the expected values in the
# expected order.
- for buffer_size in [1, 10, 100, 1000]:
- sess.run(init_op, feed_dict={buffer_size_placeholder: buffer_size})
- for i in range(100):
- self.assertEqual(i * i, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ for i in range(100):
+ self.assertEqual(i * i, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
- # We can indirectly observe that varying the buffer size has the
- # intended effect by observing when `ev` is set (on the 6th
- # invocation of `_map_py_func()`).
- # NOTE(mrry): We do not test with `buffer_size ==
- # set_event_during_invocation`, because we must consume at least
- # one element to start the prefetching.
- for buffer_size in range(1, set_event_during_invocation):
- event_will_be_set_after_consuming = (
- set_event_during_invocation - buffer_size + 1)
+ for buffer_size in [1, 10, 100, 1000]:
+ do_test(buffer_size)
- ev.clear()
- sess.run(init_op, feed_dict={buffer_size_placeholder: buffer_size})
- for i in range(event_will_be_set_after_consuming):
- self.assertFalse(ev.is_set())
- self.assertEqual(i * i, sess.run(get_next))
- ev.wait()
- for i in range(event_will_be_set_after_consuming, 100):
- self.assertEqual(i * i, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ # We can indirectly observe that varying the buffer size has the
+ # intended effect by observing when `ev` is set (on the 6th
+ # invocation of `_map_py_func()`).
+ # NOTE(mrry): We do not test with `buffer_size ==
+ # set_event_during_invocation`, because we must consume at least
+ # one element to start the prefetching.
+ def do_test_ev(buffer_size):
+ dataset = dataset_ops.Dataset.range(100).map(_map_fn).prefetch(
+ buffer_size)
+
+ get_next = self.getNext(dataset)
+
+ event_will_be_set_after_consuming = (
+ set_event_during_invocation - buffer_size + 1)
+
+ ev.clear()
+ for i in range(event_will_be_set_after_consuming):
+ self.assertFalse(ev.is_set())
+ self.assertEqual(i * i, self.evaluate(get_next()))
+ ev.wait()
+ for i in range(event_will_be_set_after_consuming, 100):
+ self.assertEqual(i * i, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
+
+ for buffer_size in range(1, set_event_during_invocation):
+ do_test_ev(buffer_size)
def testReturnList(self):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10)
- .map(lambda x: [x, constant_op.constant(37.0)]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(10):
- self.assertEqual((i, 37.0), sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(
+ 10).map(lambda x: [x, constant_op.constant(37.0)])
+ self.assertDatasetProduces(
+ dataset, expected_output=[(i, 37.0) for i in range(10)])
def testMultiOutputPyFunc(self):
# The `tf.py_func()` op returns a list of tensors for its outputs.
@@ -776,17 +714,9 @@
return script_ops.py_func(
_map_py_func, [x_tensor], [dtypes.int64, dtypes.float64])
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10).map(_map_fn))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(10):
- self.assertEqual((i, 37.0), sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(10).map(_map_fn)
+ self.assertDatasetProduces(
+ dataset, expected_output=[(i, 37.0) for i in range(10)])
def testSparse(self):
@@ -796,19 +726,9 @@
values=(i * np.array([1])),
dense_shape=np.array([1, 1]))
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10).map(_sparse))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(10):
- actual = sess.run(get_next)
- self.assertIsInstance(actual, sparse_tensor.SparseTensorValue)
- self.assertSparseValuesEqual(actual, _sparse(i))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(10).map(_sparse)
+ self.assertDatasetProduces(
+ dataset, expected_output=[_sparse(i) for i in range(10)])
def testSparseChain(self):
@@ -822,19 +742,11 @@
self.assertTrue(sparse_tensor.is_sparse(i))
return sparse_ops.sparse_concat(0, [i, i])
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10).map(_sparse).map(_check))
- init_op = iterator.initializer
- get_next = iterator.get_next()
+ dataset = dataset_ops.Dataset.range(10).map(_sparse).map(_check)
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(10):
- actual = sess.run(get_next)
- self.assertIsInstance(actual, sparse_tensor.SparseTensorValue)
- self.assertSparseValuesEqual(actual, _check(_sparse(i)).eval())
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ self.assertDatasetProduces(
+ dataset,
+ expected_output=[self.evaluate(_check(_sparse(i))) for i in range(10)])
def testParallelMapOutOfRangeError(self):
def raising_py_func(i):
@@ -843,32 +755,18 @@
else:
return i
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(105)
- .map(lambda x: script_ops.py_func(raising_py_func, [x], dtypes.int64),
- num_parallel_calls=2))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(100):
- self.assertEqual(i, sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(105).map(
+ lambda x: script_ops.py_func(raising_py_func, [x], dtypes.int64),
+ num_parallel_calls=2)
+ get_next = self.getNext(dataset)
+ for i in range(100):
+ self.assertEqual(i, self.evaluate(get_next()))
+ with self.assertRaises(errors.OutOfRangeError):
+ self.evaluate(get_next())
def testConstantOutput(self):
- iterator = dataset_ops.make_initializable_iterator(
- dataset_ops.Dataset.range(10).map(lambda x: [x, "hello", 10]))
- init_op = iterator.initializer
- get_next = iterator.get_next()
-
- with self.cached_session() as sess:
- sess.run(init_op)
- for i in range(10):
- self.assertEqual((i, b"hello", 10), sess.run(get_next))
- with self.assertRaises(errors.OutOfRangeError):
- sess.run(get_next)
+ dataset = dataset_ops.Dataset.range(10).map(lambda x: [x, "hello", 10])
+ self.assertDatasetProduces(dataset, [(i, b"hello", 10) for i in range(10)])
def testWarnOnLookupTable(self):
def collecting_function(x):
@@ -897,7 +795,7 @@
dataset_ops.Dataset.from_tensor_slices).map(
lambda ds: ds.batch(3)).flat_map(lambda x: x)
- self.assertDatasetProduces(dataset, [[1.0, 2.0, 3.0]])
+ self.assertDatasetProduces(dataset, expected_output=[[1.0, 2.0, 3.0]])
def testReturnValueError(self):
dataset = dataset_ops.Dataset.from_tensors([1.0, 2.0, 3.0])
@@ -930,11 +828,8 @@
return const_tensor
dataset = dataset.map(broken_function)
- iterator = dataset_ops.make_initializable_iterator(dataset)
-
- with self.cached_session() as sess:
- with self.assertRaisesRegexp(errors.InvalidArgumentError, "BrokenConst"):
- sess.run(iterator.initializer)
+ self.assertDatasetProduces(
+ dataset, expected_error=(errors.InvalidArgumentError, "BrokenConst"))
# pylint: disable=g-long-lambda
@parameterized.named_parameters(
@@ -957,12 +852,10 @@
return tids
dataset = make_dataset_fn(dataset, _map_fn)
- iterator = dataset_ops.make_one_shot_iterator(dataset)
- get_next = iterator.get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- tids = sess.run(get_next)
- self.assertTrue(all(tids[0] == tid for tid in tids))
+ tids = self.evaluate(get_next())
+ self.assertTrue(all(tids[0] == tid for tid in tids))
# pylint: enable=g-long-lambda
@parameterized.named_parameters(
@@ -978,30 +871,28 @@
def testShortCircuit(self, structure, map_fn, num_parallel_calls):
dataset = self.structuredDataset(structure).repeat().map(
map_fn, num_parallel_calls=num_parallel_calls)
- get_next = dataset_ops.make_one_shot_iterator(dataset).get_next()
+ get_next = self.getNext(dataset)
- with self.cached_session() as sess:
- if isinstance(structure, tuple):
- expected = map_fn(*sess.run(self.structuredElement(structure)))
- else:
- expected = map_fn(sess.run(self.structuredElement(structure)))
- self.assertEqual(expected, sess.run(get_next))
+ if isinstance(structure, tuple):
+ expected = map_fn(*self.evaluate(self.structuredElement(structure)))
+ else:
+ expected = map_fn(self.evaluate(self.structuredElement(structure)))
+ self.assertEqual(expected, self.evaluate(get_next()))
@parameterized.named_parameters(
("Sequential", None),
("Parallel", 10),
)
def testShortCircuitCapturedInput(self, num_parallel_calls):
- captured_t = array_ops.placeholder(dtypes.int64, shape=[])
+ captured_t = variables.Variable(42)
dataset = self.structuredDataset(None).repeat().map(
lambda x: captured_t, num_parallel_calls=num_parallel_calls)
- iterator = dataset_ops.make_initializable_iterator(dataset)
- get_next = iterator.get_next()
+ self.evaluate(variables.global_variables_initializer())
+ get_next = self.getNext(dataset, requires_initialization=True)
- with self.cached_session() as sess:
- sess.run(iterator.initializer, feed_dict={captured_t: 42})
- self.assertEqual(42, sess.run(get_next))
+ self.assertEqual(42, self.evaluate(get_next()))
+ # TODO(b/117581999): Add eager coverage.
@parameterized.named_parameters(
("1", 1, 1),
("2", 10, 1),
@@ -1010,7 +901,9 @@
("5", 100, 10),
("6", 100, 100),
)
- def testSloppyInterleaveInOrder(self, num_elements, num_parallel_calls):
+ @test_util.run_v1_only("b/120545219")
+ def testSkipEagerSloppyInterleaveInOrder(self, num_elements,
+ num_parallel_calls):
get_next, coordination_events = _make_coordinated_sloppy_dataset(
num_elements, num_parallel_calls)
config = config_pb2.ConfigProto(
@@ -1023,12 +916,15 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
+ # TODO(b/117581999): Add eager coverage.
@parameterized.named_parameters(
("1", 10, 10),
("2", 100, 10),
("3", 100, 100),
)
- def testSloppyInterleaveOutOfOrder(self, num_elements, num_parallel_calls):
+ @test_util.run_v1_only("b/120545219")
+ def testSkipEagerSloppyInterleaveOutOfOrder(self, num_elements,
+ num_parallel_calls):
get_next, coordination_events = _make_coordinated_sloppy_dataset(
num_elements, num_parallel_calls)
config = config_pb2.ConfigProto(
@@ -1045,6 +941,22 @@
with self.assertRaises(errors.OutOfRangeError):
sess.run(get_next)
+ @parameterized.named_parameters(
+ ("Map", None),
+ ("ParallelMap", 12),
+ )
+ def testPreserveCardinality(self, num_parallel_calls):
+
+ def py_fn(_):
+ raise StopIteration()
+
+ dataset = dataset_ops.DatasetV2.from_tensors(0).map(
+ lambda x: script_ops.py_func(py_fn, [x], dtypes.int64),
+ num_parallel_calls=num_parallel_calls)
+ get_next = self.getNext(dataset)
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(get_next())
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py b/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py
index 622ebb5..433ea62 100644
--- a/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py
+++ b/tensorflow/python/data/kernel_tests/multi_device_iterator_test.py
@@ -20,7 +20,6 @@
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.data.experimental.ops import optimization
-from tensorflow.python.data.experimental.ops.optimization_options import OptimizationOptions
from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import multi_device_iterator_ops
@@ -35,6 +34,7 @@
# TODO(b/117581999): Add eager coverage.
class MultiDeviceIteratorTest(test_base.DatasetTestBase):
+ @test_util.run_v1_only("b/120545219")
def testNoGetNext(self):
dataset = dataset_ops.Dataset.range(10)
multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
@@ -44,6 +44,7 @@
with self.test_session(config=config) as sess:
self.evaluate(multi_device_iterator.initializer)
+ @test_util.run_v1_only("b/120545219")
def testBasic(self):
dataset = dataset_ops.Dataset.range(10)
multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
@@ -60,6 +61,7 @@
self.evaluate(elem_on_1)
self.evaluate(elem_on_2)
+ @test_util.run_v1_only("b/120545219")
def testOneOnSameDevice(self):
with ops.device("/cpu:0"):
dataset = dataset_ops.Dataset.range(10)
@@ -77,6 +79,7 @@
self.evaluate(elem_on_1)
self.evaluate(elem_on_2)
+ @test_util.run_v1_only("b/120545219")
def testRepeatDevices(self):
with ops.device("/cpu:0"):
dataset = dataset_ops.Dataset.range(20)
@@ -99,6 +102,7 @@
self.evaluate(elem_on_3)
self.evaluate(elem_on_4)
+ @test_util.run_v1_only("b/120545219")
def testNotFullyDivisible(self):
dataset = dataset_ops.Dataset.range(9)
multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
@@ -116,6 +120,7 @@
self.evaluate(elem_on_1)
self.evaluate(elem_on_2)
+ @test_util.run_v1_only("b/120545219")
def testGetNextAsOptional(self):
dataset = dataset_ops.Dataset.range(9)
multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
@@ -149,6 +154,7 @@
with self.assertRaises(errors.InvalidArgumentError):
self.evaluate(elem_on_2_t)
+ @test_util.run_v1_only("b/120545219")
def testUneven(self):
dataset = dataset_ops.Dataset.range(10)
multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
@@ -166,6 +172,7 @@
self.evaluate(elem_on_1)
self.evaluate(elem_on_2)
+ @test_util.run_v1_only("b/120545219")
def testMultipleInitializations(self):
with ops.device("/cpu:0"):
epoch = array_ops.placeholder(dtypes.int64, shape=[])
@@ -259,6 +266,7 @@
with self.assertRaises(errors.InvalidArgumentError):
self.evaluate(elem_on_2_t)
+ @test_util.run_v1_only("b/120545219")
def testOptimization(self):
dataset = dataset_ops.Dataset.range(10)
dataset = dataset.apply(optimization.assert_next(["MemoryCacheImpl"]))
@@ -266,7 +274,6 @@
dataset = dataset.cache()
options = dataset_ops.Options()
- options.experimental_optimization = OptimizationOptions()
options.experimental_optimization.noop_elimination = True
dataset = dataset.with_options(options)
diff --git a/tensorflow/python/data/kernel_tests/optional_test.py b/tensorflow/python/data/kernel_tests/optional_test.py
index c2c62e9..ba5ee9b 100644
--- a/tensorflow/python/data/kernel_tests/optional_test.py
+++ b/tensorflow/python/data/kernel_tests/optional_test.py
@@ -25,6 +25,7 @@
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import optional_ops
from tensorflow.python.data.util import structure
+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 errors
@@ -361,6 +362,25 @@
with self.assertRaises(errors.InvalidArgumentError):
sess.run(elem_value_t)
+ def testFunctionBoundaries(self):
+ @def_function.function
+ def get_optional():
+ x = constant_op.constant(1.0)
+ opt = optional_ops.Optional.from_value(x)
+ # TODO(skyewm): support returning Optionals from functions?
+ return opt._variant_tensor
+
+ # TODO(skyewm): support Optional arguments?
+ @def_function.function
+ def consume_optional(opt_tensor):
+ value_structure = structure.TensorStructure(dtypes.float32, [])
+ opt = optional_ops._OptionalImpl(opt_tensor, value_structure)
+ return opt.get_value()
+
+ opt_tensor = get_optional()
+ val = consume_optional(opt_tensor)
+ self.assertEqual(self.evaluate(val), 1.0)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/data/kernel_tests/options_test.py b/tensorflow/python/data/kernel_tests/options_test.py
new file mode 100644
index 0000000..f5bad3e
--- /dev/null
+++ b/tensorflow/python/data/kernel_tests/options_test.py
@@ -0,0 +1,96 @@
+# 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 `tf.data.Options`."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.experimental.ops import optimization_options
+from tensorflow.python.data.experimental.ops import stats_options
+from tensorflow.python.data.experimental.ops import threading_options
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.platform import test
+
+
+class OptionsTest(test_base.DatasetTestBase):
+
+ def testOptionsDefault(self):
+ ds = dataset_ops.Dataset.range(0)
+ self.assertEqual(dataset_ops.Options(), ds.options())
+
+ def testOptionsOnce(self):
+ options = dataset_ops.Options()
+ ds = dataset_ops.Dataset.range(0).with_options(options).cache()
+ self.assertEqual(options, ds.options())
+
+ def testOptionsTwiceSame(self):
+ options = dataset_ops.Options()
+ options.experimental_autotune = True
+ ds = dataset_ops.Dataset.range(0).with_options(options).with_options(
+ options)
+ self.assertEqual(options, ds.options())
+
+ def testOptionsTwiceDifferent(self):
+ options1 = dataset_ops.Options()
+ options1.experimental_autotune = True
+ options2 = dataset_ops.Options()
+ options2.experimental_deterministic = False
+ ds = dataset_ops.Dataset.range(0).with_options(options1).with_options(
+ options2)
+ self.assertTrue(ds.options().experimental_autotune)
+ # Explicitly check that flag is False since assertFalse allows None
+ self.assertIs(ds.options().experimental_deterministic, False)
+
+ def testOptionsTwiceDifferentError(self):
+ options1 = dataset_ops.Options()
+ options1.experimental_autotune = True
+ options2 = dataset_ops.Options()
+ options2.experimental_autotune = False
+ with self.assertRaisesRegexp(ValueError,
+ "Cannot merge incompatible values"):
+ dataset_ops.Dataset.range(0).with_options(options1).with_options(options2)
+
+ def testOptionsMergeOptionsFromMultipleInputs(self):
+ options1 = dataset_ops.Options()
+ options1.experimental_autotune = True
+ options2 = dataset_ops.Options()
+ options2.experimental_deterministic = True
+ ds = dataset_ops.Dataset.zip(
+ (dataset_ops.Dataset.range(0).with_options(options1),
+ dataset_ops.Dataset.range(0).with_options(options2)))
+ self.assertTrue(ds.options().experimental_autotune)
+ self.assertTrue(ds.options().experimental_deterministic)
+
+ def testOptionsHaveDefaults(self):
+ options1 = dataset_ops.Options()
+ options2 = dataset_ops.Options()
+ self.assertIsNot(options1.experimental_optimization,
+ options2.experimental_optimization)
+ self.assertIsNot(options1.experimental_stats,
+ options2.experimental_stats)
+ self.assertIsNot(options1.experimental_threading,
+ options2.experimental_threading)
+ self.assertEquals(options1.experimental_optimization,
+ optimization_options.OptimizationOptions())
+ self.assertEquals(options1.experimental_stats,
+ stats_options.StatsOptions())
+ self.assertEquals(options1.experimental_threading,
+ threading_options.ThreadingOptions())
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/data/kernel_tests/test_base.py b/tensorflow/python/data/kernel_tests/test_base.py
index 85f6c9d..8793fd3 100644
--- a/tensorflow/python/data/kernel_tests/test_base.py
+++ b/tensorflow/python/data/kernel_tests/test_base.py
@@ -88,6 +88,7 @@
def assertDatasetProduces(self,
dataset,
expected_output=None,
+ expected_shapes=None,
expected_error=None,
requires_initialization=False,
num_test_iterations=1,
@@ -98,6 +99,8 @@
dataset: A dataset to check for the expected output / error.
expected_output: A list of elements that the dataset is expected to
produce.
+ expected_shapes: A list of TensorShapes which is expected to match
+ output_shapes of dataset.
expected_error: A tuple `(type, predicate)` identifying the expected error
`dataset` should raise. The `type` should match the expected exception
type, while `predicate` should either be 1) a unary function that inputs
@@ -126,6 +129,8 @@
dataset, requires_initialization=requires_initialization)
self.evaluate(get_next())
return
+ if expected_shapes:
+ self.assertEqual(expected_shapes, dataset.output_shapes)
self.assertGreater(num_test_iterations, 0)
for _ in range(num_test_iterations):
get_next = self.getNext(
diff --git a/tensorflow/python/data/ops/BUILD b/tensorflow/python/data/ops/BUILD
index 0c5acda..fbff7df 100644
--- a/tensorflow/python/data/ops/BUILD
+++ b/tensorflow/python/data/ops/BUILD
@@ -52,6 +52,7 @@
"//tensorflow/python:tensor_shape",
"//tensorflow/python/compat",
"//tensorflow/python/data/util:convert",
+ "//tensorflow/python/data/util:structure",
],
)
@@ -86,6 +87,7 @@
"//tensorflow/python:framework_ops",
"//tensorflow/python:resource_variable_ops",
"//tensorflow/python:tensor_shape",
+ "//tensorflow/python:util",
"//tensorflow/python/data/util:structure",
],
)
diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py
index ba6d203..6582ac2 100644
--- a/tensorflow/python/data/ops/dataset_ops.py
+++ b/tensorflow/python/data/ops/dataset_ops.py
@@ -187,13 +187,22 @@
RuntimeError: If eager execution is not enabled.
"""
if context.executing_eagerly():
- dataset = self._apply_options()
- return iterator_ops.EagerIterator(dataset)
+ return iterator_ops.EagerIterator(self)
else:
raise RuntimeError("dataset.__iter__() is only supported when eager "
"execution is enabled.")
@abc.abstractproperty
+ def _element_structure(self):
+ """The structure of an element of this dataset.
+
+ Returns:
+ A `Structure` object representing the structure of an element of this
+ dataset.
+ """
+ raise NotImplementedError("Dataset._element_structure")
+
+ @property
def output_classes(self):
"""Returns the class of each component of an element of this dataset.
@@ -203,9 +212,9 @@
A nested structure of Python `type` objects corresponding to each
component of an element of this dataset.
"""
- raise NotImplementedError("Dataset.output_classes")
+ return self._element_structure._to_legacy_output_classes() # pylint: disable=protected-access
- @abc.abstractproperty
+ @property
def output_shapes(self):
"""Returns the shape of each component of an element of this dataset.
@@ -213,9 +222,9 @@
A nested structure of `tf.TensorShape` objects corresponding to each
component of an element of this dataset.
"""
- raise NotImplementedError("Dataset.output_shapes")
+ return self._element_structure._to_legacy_output_shapes() # pylint: disable=protected-access
- @abc.abstractproperty
+ @property
def output_types(self):
"""Returns the type of each component of an element of this dataset.
@@ -223,7 +232,7 @@
A nested structure of `tf.DType` objects corresponding to each component
of an element of this dataset.
"""
- raise NotImplementedError("Dataset.output_types")
+ return self._element_structure._to_legacy_output_types() # pylint: disable=protected-access
def __repr__(self):
output_shapes = nest.map_structure(str, self.output_shapes)
@@ -938,15 +947,18 @@
`self.output_types`) to another nested structure of tensors.
num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
representing the number elements to process in parallel. If not
- specified, elements will be processed sequentially.
+ specified, elements will be processed sequentially. If the value
+ `tf.data.experimental.AUTOTUNE` is used, then the number of parallel
+ calls is set dynamically based on available CPU.
Returns:
Dataset: A `Dataset`.
"""
if num_parallel_calls is None:
- return MapDataset(self, map_func)
+ return MapDataset(self, map_func, preserve_cardinality=True)
else:
- return ParallelMapDataset(self, map_func, num_parallel_calls)
+ return ParallelMapDataset(
+ self, map_func, num_parallel_calls, preserve_cardinality=True)
def flat_map(self, map_func):
"""Maps `map_func` across this dataset and flattens the result.
@@ -1048,7 +1060,9 @@
num_parallel_calls: (Optional.) If specified, the implementation creates
a threadpool, which is used to fetch inputs from cycle elements
asynchronously and in parallel. The default behavior is to fetch inputs
- from cycle elements synchronously with no parallelism.
+ from cycle elements synchronously with no parallelism. If the value
+ `tf.data.experimental.AUTOTUNE` is used, then the number of parallel
+ calls is set dynamically based on available CPU.
Returns:
Dataset: A `Dataset`.
@@ -1181,27 +1195,23 @@
# Compute initial values for the state classes, shapes and types based on
# the initial state.
- state_classes = sparse.get_classes(initial_state)
- state_shapes = nest.pack_sequence_as(
- initial_state, [t.get_shape() for t in nest.flatten(initial_state)])
- state_types = nest.pack_sequence_as(
- initial_state, [t.dtype for t in nest.flatten(initial_state)])
+ state_structure = structure_lib.Structure.from_value(initial_state)
# Iteratively rerun the reduce function until reaching a fixed point on
- # `self._state_shapes`.
+ # `state_structure`.
need_to_rerun = True
while need_to_rerun:
wrapped_func = StructuredFunctionWrapper(
reduce_func,
"reduce()",
- input_classes=(state_classes, self.output_classes),
- input_shapes=(state_shapes, self.output_shapes),
- input_types=(state_types, self.output_types),
+ input_structure=structure_lib.NestedStructure(
+ (state_structure, self._element_structure)),
add_to_graph=False)
# Extract and validate class information from the returned values.
output_classes = wrapped_func.output_classes
+ state_classes = state_structure._to_legacy_output_classes() # pylint: disable=protected-access
for new_state_class, state_class in zip(
nest.flatten(output_classes), nest.flatten(state_classes)):
if not issubclass(new_state_class, state_class):
@@ -1212,6 +1222,7 @@
# Extract and validate type information from the returned values.
output_types = wrapped_func.output_types
+ state_types = state_structure._to_legacy_output_types() # pylint: disable=protected-access
for new_state_type, state_type in zip(
nest.flatten(output_types), nest.flatten(state_types)):
if new_state_type != state_type:
@@ -1222,6 +1233,7 @@
# Extract shape information from the returned values.
output_shapes = wrapped_func.output_shapes
+ state_shapes = state_structure._to_legacy_output_shapes() # pylint: disable=protected-access
flat_state_shapes = nest.flatten(state_shapes)
flat_new_state_shapes = nest.flatten(output_shapes)
weakened_state_shapes = [
@@ -1239,27 +1251,26 @@
break
if need_to_rerun:
- state_shapes = nest.pack_sequence_as(state_shapes,
- weakened_state_shapes)
+ # TODO(b/110122868): Support a "most specific compatible structure"
+ # method for combining structures, to avoid using legacy structures
+ # here.
+ state_structure = structure_lib.convert_legacy_structure(
+ state_types,
+ nest.pack_sequence_as(state_shapes, weakened_state_shapes),
+ state_classes)
reduce_func = wrapped_func.function
reduce_func.add_to_graph(ops.get_default_graph())
- return sparse.deserialize_sparse_tensors(
- nest.pack_sequence_as(
- output_types,
- gen_dataset_ops.reduce_dataset(
- self._as_variant_tensor(), # pylint: disable=protected-access
- nest.flatten(sparse.serialize_sparse_tensors(initial_state)),
- reduce_func.captured_inputs,
- f=reduce_func,
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)),
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)))),
- output_types,
- output_shapes,
- output_classes)
+ # pylint: disable=protected-access
+ return state_structure._from_compatible_tensor_list(
+ gen_dataset_ops.reduce_dataset(
+ self._as_variant_tensor(),
+ state_structure._to_tensor_list(initial_state),
+ reduce_func.captured_inputs,
+ f=reduce_func,
+ output_shapes=state_structure._flat_shapes,
+ output_types=state_structure._flat_types))
def with_options(self, options):
"""Returns a new `tf.data.Dataset` with the given options set.
@@ -1307,8 +1318,7 @@
An `Iterator` over the elements of this dataset.
"""
if context.executing_eagerly():
- dataset = self._apply_options()
- return iterator_ops.EagerIterator(dataset)
+ return iterator_ops.EagerIterator(self)
graph_level_seed, op_level_seed = core_random_seed.get_seed(None)
@@ -1396,6 +1406,13 @@
dataset.output_types, dataset.output_shapes,
dataset.output_classes)
+ @property
+ def _element_structure(self):
+ # TODO(b/110122868): Remove this override once all `Dataset` instances
+ # implement `element_structure`.
+ return structure_lib.convert_legacy_structure(
+ self.output_types, self.output_shapes, self.output_classes)
+
@staticmethod
@functools.wraps(DatasetV2.from_tensors)
def from_tensors(tensors):
@@ -1540,8 +1557,13 @@
@functools.wraps(DatasetV2.map)
def map(self, map_func, num_parallel_calls=None):
- return DatasetV1Adapter(super(DatasetV1, self).map(
- map_func, num_parallel_calls))
+ if num_parallel_calls is None:
+ return DatasetV1Adapter(
+ MapDataset(self, map_func, preserve_cardinality=False))
+ else:
+ return DatasetV1Adapter(
+ ParallelMapDataset(
+ self, map_func, num_parallel_calls, preserve_cardinality=False))
@functools.wraps(DatasetV2.flat_map)
def flat_map(self, map_func):
@@ -1599,16 +1621,8 @@
return self._dataset.options()
@property
- def output_classes(self):
- return self._dataset.output_classes
-
- @property
- def output_shapes(self):
- return self._dataset.output_shapes
-
- @property
- def output_types(self):
- return self._dataset.output_types
+ def _element_structure(self):
+ return self._dataset._element_structure # pylint: disable=protected-access
def __iter__(self):
return iter(self._dataset)
@@ -1687,8 +1701,8 @@
name="experimental_deterministic",
ty=bool,
docstring=
- "Whether to dynamically adjust the values of tunable parameters (e.g. "
- "degrees of parallelism).")
+ "Whether the outputs need to be produced in deterministic order."
+ )
experimental_numa_aware = options_lib.create_option(
name="experimental_numa_aware",
@@ -1698,38 +1712,27 @@
experimental_optimization = options_lib.create_option(
name="experimental_optimization",
ty=optimization_options.OptimizationOptions,
- docstring="Associates the given optimization options with the dataset.")
+ docstring="Associates the given optimization options with the dataset.",
+ default_factory=optimization_options.OptimizationOptions)
experimental_stats = options_lib.create_option(
name="experimental_stats",
ty=stats_options.StatsOptions,
- docstring="Associates the given statistics options with the dataset.")
+ docstring="Associates the given statistics options with the dataset.",
+ default_factory=stats_options.StatsOptions)
experimental_threading = options_lib.create_option(
name="experimental_threading",
ty=threading_options.ThreadingOptions,
- docstring="Associates the given threading options with the dataset.")
+ docstring="Associates the given threading options with the dataset.",
+ default_factory=threading_options.ThreadingOptions)
def _static_optimizations(self):
"""Produces the list of enabled static optimizations."""
result = []
- exp_optimization_options = self.experimental_optimization
- if exp_optimization_options:
- optimizations = [
- "filter_fusion",
- "hoist_random_uniform",
- "map_and_batch_fusion",
- "map_and_filter_fusion",
- "map_fusion",
- "map_parallelization",
- "map_vectorization",
- "noop_elimination",
- "shuffle_and_repeat_fusion",
- ]
- for optimization in optimizations:
- if getattr(exp_optimization_options, optimization):
- result.append(optimization)
+ result.extend(self.experimental_optimization._static_optimizations()) # pylint: disable=protected-access
+
if self.experimental_numa_aware:
result.append("make_numa_aware")
if self.experimental_deterministic is False:
@@ -1780,16 +1783,8 @@
"""Represents a unary dataset with the same input and output structure."""
@property
- def output_classes(self):
- return self._input_dataset.output_classes # pylint: disable=protected-access
-
- @property
- def output_shapes(self):
- return self._input_dataset.output_shapes # pylint: disable=protected-access
-
- @property
- def output_types(self):
- return self._input_dataset.output_types # pylint: disable=protected-access
+ def _element_structure(self):
+ return self._input_dataset._element_structure # pylint: disable=protected-access
class TensorDataset(DatasetSource):
@@ -1805,31 +1800,16 @@
t, name="component_%d" % i)
for i, t in enumerate(nest.flatten(tensors))
])
-
- self._tensors = sparse.serialize_sparse_tensors(tensors)
- self._output_classes = sparse.get_classes(tensors)
- self._output_shapes = nest.pack_sequence_as(
- tensors, [t.get_shape() for t in nest.flatten(tensors)])
- self._output_types = nest.pack_sequence_as(
- tensors, [t.dtype for t in nest.flatten(tensors)])
+ self._structure = structure_lib.Structure.from_value(tensors)
+ self._tensors = self._structure._to_tensor_list(tensors) # pylint: disable=protected-access
def _as_variant_tensor(self):
return gen_dataset_ops.tensor_dataset(
- nest.flatten(self._tensors),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
+ self._tensors, output_shapes=self._structure._flat_shapes) # pylint: disable=protected-access
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
class TensorSliceDataset(DatasetSource):
@@ -1845,37 +1825,26 @@
t, name="component_%d" % i)
for i, t in enumerate(nest.flatten(tensors))
])
- flat_tensors = nest.flatten(tensors)
+
+ batched_structure = structure_lib.Structure.from_value(tensors)
+ # pylint: disable=protected-access
+ self._tensors = batched_structure._to_batched_tensor_list(tensors)
+ self._structure = batched_structure._unbatch()
+ # pylint: enable=protected-access
batch_dim = tensor_shape.Dimension(tensor_shape.dimension_value(
- flat_tensors[0].get_shape()[0]))
- for t in flat_tensors[1:]:
+ self._tensors[0].get_shape()[0]))
+ for t in self._tensors[1:]:
batch_dim.assert_is_compatible_with(tensor_shape.Dimension(
tensor_shape.dimension_value(t.get_shape()[0])))
- self._tensors = sparse.serialize_many_sparse_tensors(tensors)
- self._output_classes = sparse.get_classes(tensors)
- self._output_shapes = nest.pack_sequence_as(
- tensors, [t.get_shape()[1:] for t in nest.flatten(tensors)])
- self._output_types = nest.pack_sequence_as(
- tensors, [t.dtype for t in nest.flatten(tensors)])
def _as_variant_tensor(self):
return gen_dataset_ops.tensor_slice_dataset(
- nest.flatten(self._tensors),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
+ self._tensors, output_shapes=self._structure._flat_shapes) # pylint: disable=protected-access
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
class SparseTensorSliceDataset(DatasetSource):
@@ -1888,28 +1857,22 @@
raise TypeError("`sparse_tensor` must be a `tf.SparseTensor` object.")
self._sparse_tensor = sparse_tensor
+ indices_shape = self._sparse_tensor.indices.get_shape()
+ shape_shape = self._sparse_tensor.dense_shape.get_shape()
+ rank = (indices_shape.dims[1] - 1).merge_with(shape_shape.dims[0] - 1)
+ self._structure = structure_lib.NestedStructure(
+ (structure_lib.TensorStructure(dtypes.int64, [None, rank]),
+ structure_lib.TensorStructure(self._sparse_tensor.dtype, [None]),
+ structure_lib.TensorStructure(dtypes.int64, [rank])))
+
def _as_variant_tensor(self):
return gen_dataset_ops.sparse_tensor_slice_dataset(
self._sparse_tensor.indices, self._sparse_tensor.values,
self._sparse_tensor.dense_shape)
@property
- def output_classes(self):
- return (ops.Tensor, ops.Tensor, ops.Tensor)
-
- @property
- def output_shapes(self):
- indices_shape = self._sparse_tensor.indices.get_shape()
- shape_shape = self._sparse_tensor.dense_shape.get_shape()
- rank = (indices_shape.dims[1] - 1).merge_with(shape_shape.dims[0] - 1)
- num_values = tensor_shape.Dimension(None)
- return (tensor_shape.TensorShape([num_values, rank]),
- tensor_shape.TensorShape([num_values]),
- tensor_shape.TensorShape([rank]))
-
- @property
- def output_types(self):
- return (dtypes.int64, self._sparse_tensor.dtype, dtypes.int64)
+ def _element_structure(self):
+ return self._structure
class _VariantDataset(DatasetV2):
@@ -1927,18 +1890,11 @@
return []
@property
- def output_classes(self):
- return self._structure._to_legacy_output_classes() # pylint: disable=protected-access
-
- @property
- def output_shapes(self):
- return self._structure._to_legacy_output_shapes() # pylint: disable=protected-access
-
- @property
- def output_types(self):
- return self._structure._to_legacy_output_types() # pylint: disable=protected-access
+ def _element_structure(self):
+ return self._structure
+@tf_export("data.experimental.DatasetStructure")
class DatasetStructure(structure_lib.Structure):
"""Represents a `Dataset` of structured values."""
@@ -1962,6 +1918,9 @@
def _to_tensor_list(self, value):
return [value._as_variant_tensor()] # pylint: disable=protected-access
+ def _to_batched_tensor_list(self, value):
+ raise NotImplementedError("Unbatching for `tf.data.Dataset` objects.")
+
def _from_tensor_list(self, flat_value):
if (len(flat_value) != 1 or flat_value[0].dtype != dtypes.variant or
not flat_value[0].shape.is_compatible_with(tensor_shape.scalar())):
@@ -1975,11 +1934,7 @@
@staticmethod
def from_value(value):
- # TODO(b/110122868): We can simplify this when a `Dataset` object has a
- # `Structure`-valued property.
- element_structure = structure_lib.Structure._from_legacy_structure(
- value.output_types, value.output_shapes, value.output_classes)
- return DatasetStructure(element_structure)
+ return DatasetStructure(value._element_structure) # pylint: disable=protected-access
def _to_legacy_output_types(self):
return self
@@ -1993,6 +1948,9 @@
def _batch(self, batch_size):
raise NotImplementedError("Batching for `tf.data.Dataset` objects.")
+ def _unbatch(self):
+ raise NotImplementedError("Unbatching for `tf.data.Dataset` objects.")
+
# pylint: disable=protected-access
structure_lib.Structure._register_custom_converter(DatasetV2,
@@ -2011,6 +1969,7 @@
input_classes=None,
input_shapes=None,
input_types=None,
+ input_structure=None,
add_to_graph=True,
defun_kwargs=None):
"""Creates a new `StructuredFunctionWrapper` for the given function.
@@ -2029,6 +1988,8 @@
arguments.
input_types: (Optional.) A nested structure of `tf.DType`. If given, this
argument defines the element types and structure for `func` arguments.
+ input_structure: (Optional.) A `Structure` object. If given, this argument
+ defines the element types and structure for `func` arguments.
add_to_graph: (Optional.) If `True`, the function will be added to the
default graph.
defun_kwargs: (Optional.) A dictionary mapping string argument names to
@@ -2039,24 +2000,28 @@
ValueError: If an invalid combination of `dataset`, `input_classes`,
`input_shapes`, and `input_types` is passed.
"""
- if dataset is None:
- if input_classes is None or input_shapes is None or input_types is None:
- raise ValueError("Either `dataset`, or all of `input_classes`, "
- "`input_shapes`, and `input_types` must be specified.")
- self._input_shapes = input_shapes
- self._input_types = input_types
- self._input_classes = input_classes
+ if input_structure is None:
+ if dataset is None:
+ if input_classes is None or input_shapes is None or input_types is None:
+ raise ValueError("Either `dataset`, `input_structure` or all of "
+ "`input_classes`, `input_shapes`, and `input_types` "
+ "must be specified.")
+ self._input_structure = structure_lib.convert_legacy_structure(
+ input_types, input_shapes, input_classes)
+ else:
+ if not (input_classes is None and input_shapes is None and
+ input_types is None):
+ raise ValueError("Either `dataset`, `input_structure` or all of "
+ "`input_classes`, `input_shapes`, and `input_types` "
+ "must be specified.")
+ self._input_structure = dataset._element_structure # pylint: disable=protected-access
else:
- if not (input_classes is None and input_shapes is None and
- input_types is None):
- raise ValueError("Either `dataset`, or all of `input_classes`, "
- "`input_shapes`, and `input_types` must be specified.")
- self._input_shapes = dataset.output_shapes
- self._input_types = dataset.output_types
- self._input_classes = dataset.output_classes
-
- self._input_structure = structure_lib.Structure._from_legacy_structure( # pylint: disable=protected-access
- self._input_types, self._input_shapes, self._input_classes)
+ if not (dataset is None and input_classes is None and input_shapes is None
+ and input_types is None):
+ raise ValueError("Either `dataset`, `input_structure`, or all of "
+ "`input_classes`, `input_shapes`, and `input_types` "
+ "must be specified.")
+ self._input_structure = input_structure
self._transformation_name = transformation_name
readable_transformation_name = transformation_name.replace(
@@ -2133,11 +2098,9 @@
return self._function
-def flat_structure(dataset=None, structure=None):
+def flat_structure(dataset):
"""Helper for setting `output_shapes` and `output_types` attrs of Dataset ops.
- Either `dataset` or `structure` must be passed to this function.
-
Most Dataset op constructors expect `output_shapes` and `output_types`
arguments that represent the flattened structure of an element. This helper
function generates these attrs as a keyword argument dictionary, allowing
@@ -2145,17 +2108,14 @@
`**flat_structure(self)` to the op constructor.
Args:
- dataset: (Optional.) A `tf.data.Dataset`.
- structure: (Optional.) A `Structure`.
+ dataset: A `tf.data.Dataset`.
Returns:
A dictionary of keyword arguments that can be passed to many Dataset op
constructors.
"""
# pylint: disable=protected-access
- if structure is None:
- structure = structure_lib.Structure._from_legacy_structure(
- dataset.output_types, dataset.output_shapes, dataset.output_classes)
+ structure = dataset._element_structure
return {
"output_shapes": structure._flat_shapes,
"output_types": structure._flat_types,
@@ -2181,70 +2141,39 @@
destroyed. The return value is ignored.
"""
super(_GeneratorDataset, self).__init__()
- # These members will be initialized by `tf_init_func`.
- self._state_classes = None
- self._state_shapes = None
- self._state_types = None
-
self._init_args = init_args
- init_args_classes = sparse.get_classes(init_args)
- init_args_shapes = nest.pack_sequence_as(
- init_args, [t.get_shape() for t in nest.flatten(init_args)])
- init_args_types = nest.pack_sequence_as(
- init_args, [t.dtype for t in nest.flatten(init_args)])
+ self._init_structure = structure_lib.Structure.from_value(init_args)
- wrapped_init_func = StructuredFunctionWrapper(
+ self._init_func = StructuredFunctionWrapper(
init_func,
self._transformation_name(),
- input_classes=init_args_classes,
- input_shapes=init_args_shapes,
- input_types=init_args_types)
- self._state_classes = wrapped_init_func.output_classes
- self._state_shapes = wrapped_init_func.output_shapes
- self._state_types = wrapped_init_func.output_types
- self._init_func = wrapped_init_func.function
+ input_structure=self._init_structure)
- wrapped_next_func = StructuredFunctionWrapper(
+ self._next_func = StructuredFunctionWrapper(
next_func,
self._transformation_name(),
- input_classes=self._state_classes,
- input_shapes=self._state_shapes,
- input_types=self._state_types)
- self._output_classes = wrapped_next_func.output_classes
- self._output_shapes = wrapped_next_func.output_shapes
- self._output_types = wrapped_next_func.output_types
- self._next_func = wrapped_next_func.function
+ input_structure=self._init_func.output_structure)
- wrapped_finalize_func = StructuredFunctionWrapper(
+ self._finalize_func = StructuredFunctionWrapper(
finalize_func,
self._transformation_name(),
- input_classes=self._state_classes,
- input_shapes=self._state_shapes,
- input_types=self._state_types)
- self._finalize_func = wrapped_finalize_func.function
+ input_structure=self._init_func.output_structure)
def _as_variant_tensor(self):
return gen_dataset_ops.generator_dataset(
- nest.flatten(self._init_args) + self._init_func.captured_inputs,
- self._next_func.captured_inputs,
- self._finalize_func.captured_inputs,
- init_func=self._init_func,
- next_func=self._next_func,
- finalize_func=self._finalize_func,
+ self._init_structure._to_tensor_list(self._init_args) # pylint: disable=protected-access
+ + self._init_func.function.captured_inputs,
+ self._next_func.function.captured_inputs,
+ self._finalize_func.function.captured_inputs,
+ init_func=self._init_func.function,
+ next_func=self._next_func.function,
+ finalize_func=self._finalize_func.function,
**flat_structure(self))
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._next_func.output_structure
def _transformation_name(self):
return "Dataset.from_generator()"
@@ -2267,6 +2196,10 @@
"structure of `Dataset` objects.")
raise TypeError(message)
self._datasets = datasets
+ self._structure = structure_lib.NestedStructure(
+ nest.pack_sequence_as(
+ self._datasets,
+ [ds._element_structure for ds in nest.flatten(self._datasets)])) # pylint: disable=protected-access
def _as_variant_tensor(self):
# pylint: disable=protected-access
@@ -2279,22 +2212,8 @@
return nest.flatten(self._datasets)
@property
- def output_classes(self):
- return nest.pack_sequence_as(
- self._datasets,
- [ds.output_classes for ds in nest.flatten(self._datasets)])
-
- @property
- def output_shapes(self):
- return nest.pack_sequence_as(
- self._datasets,
- [ds.output_shapes for ds in nest.flatten(self._datasets)])
-
- @property
- def output_types(self):
- return nest.pack_sequence_as(
- self._datasets,
- [ds.output_types for ds in nest.flatten(self._datasets)])
+ def _element_structure(self):
+ return self._structure
class ConcatenateDataset(DatasetV2):
@@ -2306,26 +2225,29 @@
self._input_dataset = input_dataset
self._dataset_to_concatenate = dataset_to_concatenate
- self._output_types = input_dataset.output_types
- if self._output_types != dataset_to_concatenate.output_types:
+ output_types = input_dataset.output_types
+ if output_types != dataset_to_concatenate.output_types:
raise TypeError(
"Two datasets to concatenate have different types %s and %s" %
- (self._output_types, dataset_to_concatenate.output_types))
+ (output_types, dataset_to_concatenate.output_types))
- self._output_classes = input_dataset.output_classes
- if self._output_classes != dataset_to_concatenate.output_classes:
+ output_classes = input_dataset.output_classes
+ if output_classes != dataset_to_concatenate.output_classes:
raise TypeError(
"Two datasets to concatenate have different classes %s and %s" %
- (self._output_classes, dataset_to_concatenate.output_classes))
+ (output_classes, dataset_to_concatenate.output_classes))
input_shapes = self._input_dataset.output_shapes
- self._output_shapes = nest.pack_sequence_as(input_shapes, [
+ output_shapes = nest.pack_sequence_as(input_shapes, [
ts1.most_specific_compatible_shape(ts2)
for (ts1, ts2) in zip(
nest.flatten(input_shapes),
nest.flatten(self._dataset_to_concatenate.output_shapes))
])
+ self._structure = structure_lib.convert_legacy_structure(
+ output_types, output_shapes, output_classes)
+
self._input_datasets = [input_dataset, dataset_to_concatenate]
def _as_variant_tensor(self):
@@ -2340,16 +2262,8 @@
return [self._input_dataset, self._dataset_to_concatenate]
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
class RepeatDataset(UnaryUnchangedStructureDataset):
@@ -2408,16 +2322,8 @@
**flat_structure(self))
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.int64
+ def _element_structure(self):
+ return structure_lib.TensorStructure(dtypes.int64, [])
class CacheDataset(UnaryUnchangedStructureDataset):
@@ -2530,37 +2436,26 @@
self._drop_remainder = ops.convert_to_tensor(
drop_remainder, dtype=dtypes.bool, name="drop_remainder")
- # pylint: disable=protected-access
- input_structure = structure_lib.Structure._from_legacy_structure(
- input_dataset.output_types, input_dataset.output_shapes,
- input_dataset.output_classes)
constant_drop_remainder = tensor_util.constant_value(self._drop_remainder)
+ # pylint: disable=protected-access
if constant_drop_remainder:
# NOTE(mrry): `constant_drop_remainder` may be `None` (unknown statically)
# or `False` (explicitly retaining the remainder).
- self._output_structure = input_structure._batch(
+ self._structure = input_dataset._element_structure._batch(
tensor_util.constant_value(self._batch_size))
else:
- self._output_structure = input_structure._batch(None)
+ self._structure = input_dataset._element_structure._batch(None)
def _as_variant_tensor(self):
return gen_dataset_ops.batch_dataset_v2(
self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
batch_size=self._batch_size,
drop_remainder=self._drop_remainder,
- **flat_structure(structure=self._output_structure))
+ **flat_structure(self))
@property
- def output_classes(self):
- return self._output_structure._to_legacy_output_classes() # pylint: disable=protected-access
-
- @property
- def output_shapes(self):
- return self._output_structure._to_legacy_output_shapes() # pylint: disable=protected-access
-
- @property
- def output_types(self):
- return self._output_structure._to_legacy_output_types() # pylint: disable=protected-access
+ def _element_structure(self):
+ return self._structure
def _is_padded_shape_compatible_with(padded_shape, input_component_shape):
@@ -2709,22 +2604,34 @@
self._drop_remainder = ops.convert_to_tensor(
drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+ def _padded_shape_to_batch_shape(s):
+ return tensor_shape.vector(
+ tensor_util.constant_value(self._batch_size) if smart_cond.
+ smart_constant_value(self._drop_remainder) else None).concatenate(
+ tensor_util.constant_value_as_shape(s))
+
+ output_shapes = nest.map_structure(
+ _padded_shape_to_batch_shape, self._padded_shapes)
+ self._structure = structure_lib.convert_legacy_structure(
+ self._input_dataset.output_types, output_shapes,
+ self._input_dataset.output_classes)
+
def _as_variant_tensor(self):
+ # pylint: disable=protected-access
# TODO(jsimsa): Switch to using v2 only any time after 6/30/2018.
if smart_cond.smart_constant_value(self._drop_remainder) is False:
return gen_dataset_ops.padded_batch_dataset(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._input_dataset._as_variant_tensor(),
batch_size=self._batch_size,
padded_shapes=[
ops.convert_to_tensor(s, dtype=dtypes.int64)
for s in nest.flatten(self._padded_shapes)
],
padding_values=nest.flatten(self._padding_values),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
+ output_shapes=self._structure._flat_shapes)
else:
return gen_dataset_ops.padded_batch_dataset_v2(
- self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
+ self._input_dataset._as_variant_tensor(),
batch_size=self._batch_size,
padded_shapes=[
ops.convert_to_tensor(s, dtype=dtypes.int64)
@@ -2732,27 +2639,11 @@
],
padding_values=nest.flatten(self._padding_values),
drop_remainder=self._drop_remainder,
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
+ output_shapes=self._structure._flat_shapes)
@property
- def output_classes(self):
- return self._input_dataset.output_classes
-
- @property
- def output_shapes(self):
-
- def _padded_shape_to_batch_shape(s):
- return tensor_shape.vector(
- tensor_util.constant_value(self._batch_size) if smart_cond.
- smart_constant_value(self._drop_remainder) else None).concatenate(
- tensor_util.constant_value_as_shape(s))
-
- return nest.map_structure(_padded_shape_to_batch_shape, self._padded_shapes)
-
- @property
- def output_types(self):
- return self._input_dataset.output_types
+ def _element_structure(self):
+ return self._structure
def _should_unpack_args(args):
@@ -2781,11 +2672,16 @@
class MapDataset(UnaryDataset):
"""A `Dataset` that maps a function over elements in its input."""
- def __init__(self, input_dataset, map_func, use_inter_op_parallelism=True):
+ def __init__(self,
+ input_dataset,
+ map_func,
+ use_inter_op_parallelism=True,
+ preserve_cardinality=False):
"""See `Dataset.map()` for details."""
super(MapDataset, self).__init__(input_dataset)
self._input_dataset = input_dataset
self._use_inter_op_parallelism = use_inter_op_parallelism
+ self._preserve_cardinality = preserve_cardinality
self._map_func = StructuredFunctionWrapper(
map_func, self._transformation_name(), dataset=input_dataset)
@@ -2796,22 +2692,15 @@
self._map_func.function.captured_inputs,
f=self._map_func.function,
use_inter_op_parallelism=self._use_inter_op_parallelism,
- **flat_structure(structure=self._map_func.output_structure))
+ preserve_cardinality=self._preserve_cardinality,
+ **flat_structure(self))
def _functions(self):
return [self._map_func]
@property
- def output_classes(self):
- return self._map_func.output_classes
-
- @property
- def output_shapes(self):
- return self._map_func.output_shapes
-
- @property
- def output_types(self):
- return self._map_func.output_types
+ def _element_structure(self):
+ return self._map_func.output_structure
def _transformation_name(self):
return "Dataset.map()"
@@ -2824,10 +2713,11 @@
input_dataset,
map_func,
num_parallel_calls,
- use_inter_op_parallelism=True):
+ use_inter_op_parallelism=True,
+ preserve_cardinality=False):
"""See `Dataset.map()` for details."""
- super(ParallelMapDataset, self).__init__(input_dataset, map_func,
- use_inter_op_parallelism)
+ super(ParallelMapDataset, self).__init__(
+ input_dataset, map_func, use_inter_op_parallelism, preserve_cardinality)
self._num_parallel_calls = ops.convert_to_tensor(
num_parallel_calls, dtype=dtypes.int32, name="num_parallel_calls")
@@ -2841,7 +2731,8 @@
f=self._map_func.function,
num_parallel_calls=self._num_parallel_calls,
use_inter_op_parallelism=self._use_inter_op_parallelism,
- **flat_structure(structure=self._map_func.output_structure))
+ preserve_cardinality=self._preserve_cardinality,
+ **flat_structure(self))
class FlatMapDataset(UnaryDataset):
@@ -2856,7 +2747,7 @@
map_func, self._transformation_name(), dataset=input_dataset)
if not isinstance(self._map_func.output_structure, DatasetStructure):
raise TypeError("`map_func` must return a `Dataset` object.")
- self._output_structure = self._map_func.output_structure._element_structure # pylint: disable=protected-access
+ self._structure = self._map_func.output_structure._element_structure # pylint: disable=protected-access
def _functions(self):
return [self._map_func]
@@ -2866,19 +2757,11 @@
self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access
self._map_func.function.captured_inputs,
f=self._map_func.function,
- **flat_structure(structure=self._output_structure))
+ **flat_structure(self))
@property
- def output_classes(self):
- return self._output_structure._to_legacy_output_classes() # pylint: disable=protected-access
-
- @property
- def output_shapes(self):
- return self._output_structure._to_legacy_output_shapes() # pylint: disable=protected-access
-
- @property
- def output_types(self):
- return self._output_structure._to_legacy_output_types() # pylint: disable=protected-access
+ def _element_structure(self):
+ return self._structure
def _transformation_name(self):
return "Dataset.flat_map()"
@@ -2904,7 +2787,7 @@
self._cycle_length,
self._block_length,
f=self._map_func.function,
- **flat_structure(structure=self._output_structure))
+ **flat_structure(self))
def _transformation_name(self):
return "Dataset.interleave()"
@@ -2935,7 +2818,7 @@
self._block_length,
self._num_parallel_calls,
f=self._map_func.function,
- **flat_structure(structure=self._output_structure))
+ **flat_structure(self))
def _transformation_name(self):
return "Dataset.interleave()"
@@ -2950,9 +2833,8 @@
self._input_dataset = input_dataset
wrapped_func = StructuredFunctionWrapper(
predicate, self._transformation_name(), dataset=input_dataset)
- if not (
- wrapped_func.output_types == dtypes.bool and
- wrapped_func.output_shapes.is_compatible_with(tensor_shape.scalar())):
+ if not wrapped_func.output_structure.is_compatible_with(
+ structure_lib.TensorStructure(dtypes.bool, [])):
raise ValueError("`predicate` must return a scalar boolean tensor.")
self._predicate = wrapped_func
@@ -3002,19 +2884,17 @@
stride, dtype=dtypes.int64, name="stride")
self._drop_remainder = ops.convert_to_tensor(
drop_remainder, dtype=dtypes.bool, name="drop_remainder")
- self._output_classes = nest.pack_sequence_as(
+ nest_of_structures = nest.pack_sequence_as(
input_dataset.output_classes,
[
- DatasetStructure(
- structure_lib.Structure._from_legacy_structure( # pylint: disable=protected-access
- output_type, output_shape, output_class))
+ DatasetStructure(structure_lib.convert_legacy_structure(
+ output_type, output_shape, output_class))
for output_class, output_shape, output_type in zip(
nest.flatten(input_dataset.output_classes),
nest.flatten(input_dataset.output_shapes),
nest.flatten(input_dataset.output_types))
])
- self._output_shapes = self._output_classes
- self._output_types = self._output_classes
+ self._structure = structure_lib.NestedStructure(nest_of_structures)
def _as_variant_tensor(self):
return gen_dataset_ops.window_dataset(
@@ -3026,16 +2906,8 @@
**flat_structure(self))
@property
- def output_classes(self):
- return self._output_classes
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_types(self):
- return self._output_types
+ def _element_structure(self):
+ return self._structure
class _OptionsDataset(UnaryUnchangedStructureDataset):
diff --git a/tensorflow/python/data/ops/iterator_ops.py b/tensorflow/python/data/ops/iterator_ops.py
index e2ca64c..d0e91b0 100644
--- a/tensorflow/python/data/ops/iterator_ops.py
+++ b/tensorflow/python/data/ops/iterator_ops.py
@@ -23,7 +23,6 @@
from tensorflow.python.compat import compat
from tensorflow.python.data.ops import optional_ops
from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
from tensorflow.python.data.util import structure as structure_lib
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
@@ -100,10 +99,8 @@
raise ValueError("If `structure` is not specified, all of "
"`output_types`, `output_shapes`, and `output_classes`"
" must be specified.")
- # pylint: disable=protected-access
- self._structure = structure_lib.Structure._from_legacy_structure(
+ self._structure = structure_lib.convert_legacy_structure(
output_types, output_shapes, output_classes)
- # pylint: enable=protected-access
self._string_handle = gen_dataset_ops.iterator_to_string_handle(
self._iterator_resource)
@@ -190,34 +187,32 @@
if output_classes is None:
output_classes = nest.map_structure(lambda _: ops.Tensor, output_types)
nest.assert_same_structure(output_types, output_shapes)
+ output_structure = structure_lib.convert_legacy_structure(
+ output_types, output_shapes, output_classes)
if shared_name is None:
shared_name = ""
+ # pylint: disable=protected-access
if compat.forward_compatible(2018, 8, 3):
if _device_stack_is_empty():
with ops.device("/cpu:0"):
iterator_resource = gen_dataset_ops.iterator_v2(
container="",
shared_name=shared_name,
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)))
+ output_types=output_structure._flat_types,
+ output_shapes=output_structure._flat_shapes)
else:
iterator_resource = gen_dataset_ops.iterator_v2(
container="",
shared_name=shared_name,
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)))
+ output_types=output_structure._flat_types,
+ output_shapes=output_structure._flat_shapes)
else:
iterator_resource = gen_dataset_ops.iterator(
container="",
shared_name=shared_name,
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)))
+ output_types=output_structure._flat_types,
+ output_shapes=output_structure._flat_shapes)
+ # pylint: enable=protected-access
return Iterator(iterator_resource, None, output_types, output_shapes,
output_classes)
@@ -280,30 +275,28 @@
if output_classes is None:
output_classes = nest.map_structure(lambda _: ops.Tensor, output_types)
nest.assert_same_structure(output_types, output_shapes)
+ output_structure = structure_lib.convert_legacy_structure(
+ output_types, output_shapes, output_classes)
string_handle = ops.convert_to_tensor(string_handle, dtype=dtypes.string)
+ # pylint: disable=protected-access
if compat.forward_compatible(2018, 8, 3):
if _device_stack_is_empty():
with ops.device("/cpu:0"):
iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(
string_handle,
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)))
+ output_types=output_structure._flat_types,
+ output_shapes=output_structure._flat_shapes)
else:
iterator_resource = gen_dataset_ops.iterator_from_string_handle_v2(
string_handle,
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)))
+ output_types=output_structure._flat_types,
+ output_shapes=output_structure._flat_shapes)
else:
iterator_resource = gen_dataset_ops.iterator_from_string_handle(
string_handle,
- output_types=nest.flatten(
- sparse.as_dense_types(output_types, output_classes)),
- output_shapes=nest.flatten(
- sparse.as_dense_shapes(output_shapes, output_classes)))
+ output_types=output_structure._flat_types,
+ output_shapes=output_structure._flat_shapes)
+ # pylint: enable=protected-access
return Iterator(iterator_resource, None, output_types, output_shapes,
output_classes)
@@ -530,8 +523,9 @@
self._device = context.context().device_name
with ops.device("/cpu:0"):
# pylint: disable=protected-access
+ dataset = dataset._apply_options()
ds_variant = dataset._as_variant_tensor()
- self._structure = structure_lib.Structure._from_legacy_structure(
+ self._structure = structure_lib.convert_legacy_structure(
dataset.output_types, dataset.output_shapes, dataset.output_classes)
self._flat_output_types = self._structure._flat_types
self._flat_output_shapes = self._structure._flat_shapes
@@ -543,6 +537,7 @@
# Delete the resource when this object is deleted
self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
handle=self._resource, handle_device=self._device)
+ # pylint: enable=protected-access
def __iter__(self):
return self
diff --git a/tensorflow/python/data/ops/multi_device_iterator_ops.py b/tensorflow/python/data/ops/multi_device_iterator_ops.py
index 0f9add6..45d0156 100644
--- a/tensorflow/python/data/ops/multi_device_iterator_ops.py
+++ b/tensorflow/python/data/ops/multi_device_iterator_ops.py
@@ -19,8 +19,6 @@
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
-from tensorflow.python.data.util import nest
-from tensorflow.python.data.util import sparse
from tensorflow.python.eager import context
from tensorflow.python.eager import function
from tensorflow.python.framework import dtypes
@@ -36,16 +34,9 @@
"""A `dummy` generator dataset."""
def __init__(self, shard_num, multi_device_iterator_resource, incarnation_id,
- source_device, target_device, output_shapes, output_types,
- output_classes):
+ source_device, target_device, element_structure):
self._target_device = target_device
- self._output_types = output_types
- self._output_shapes = output_shapes
- self._output_classes = output_classes
- self._flat_output_shapes = nest.flatten(
- sparse.as_dense_shapes(self._output_shapes, self._output_classes))
- self._flat_output_types = nest.flatten(
- sparse.as_dense_types(self._output_types, self._output_classes))
+ self._structure = element_structure
multi_device_iterator_string_handle = (
gen_dataset_ops.multi_device_iterator_to_string_handle(
@@ -70,17 +61,18 @@
@function.defun(input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _next_func(string_handle):
+ # pylint: disable=protected-access
multi_device_iterator = (
gen_dataset_ops.multi_device_iterator_from_string_handle(
string_handle=string_handle,
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes))
+ output_types=self._structure._flat_types,
+ output_shapes=self._structure._flat_shapes))
return gen_dataset_ops.multi_device_iterator_get_next_from_shard(
multi_device_iterator=multi_device_iterator,
shard_num=shard_num,
incarnation_id=incarnation_id,
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
+ output_types=self._structure._flat_types,
+ output_shapes=self._structure._flat_shapes)
next_func_concrete = _next_func._get_concrete_function_internal() # pylint: disable=protected-access
@@ -90,9 +82,8 @@
def _remote_next_func(string_handle):
return functional_ops.remote_call(
target=source_device,
- args=[string_handle] +
- next_func_concrete.captured_inputs,
- Tout=self._flat_output_types,
+ args=[string_handle] + next_func_concrete.captured_inputs,
+ Tout=self._structure._flat_types, # pylint: disable=protected-access
f=next_func_concrete)
self._next_func = _remote_next_func._get_concrete_function_internal() # pylint: disable=protected-access
@@ -108,8 +99,7 @@
def _remote_finalize_func(string_handle):
return functional_ops.remote_call(
target=source_device,
- args=[string_handle] +
- finalize_func_concrete.captured_inputs,
+ args=[string_handle] + finalize_func_concrete.captured_inputs,
Tout=[dtypes.int64],
f=finalize_func_concrete)
@@ -126,24 +116,15 @@
init_func=self._init_func,
next_func=self._next_func,
finalize_func=self._finalize_func,
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes)
+ **dataset_ops.flat_structure(self))
def _inputs(self):
# TODO(b/116506223): Determine which datasets should be used as inputs here.
return []
@property
- def output_types(self):
- return self._output_types
-
- @property
- def output_shapes(self):
- return self._output_shapes
-
- @property
- def output_classes(self):
- return self._output_classes
+ def _element_structure(self):
+ return self._structure
class MultiDeviceIterator(object):
@@ -183,13 +164,6 @@
self._source_device = source_device
self._source_device_tensor = ops.convert_to_tensor(source_device)
- self._flat_output_shapes = nest.flatten(
- sparse.as_dense_shapes(self._dataset.output_shapes,
- self._dataset.output_classes))
- self._flat_output_types = nest.flatten(
- sparse.as_dense_types(self._dataset.output_types,
- self._dataset.output_classes))
-
# Create the MultiDeviceIterator.
with ops.device(self._source_device):
self._multi_device_iterator_resource = (
@@ -197,8 +171,7 @@
devices=self._devices,
shared_name="",
container="",
- output_types=self._flat_output_types,
- output_shapes=self._flat_output_shapes))
+ **dataset_ops.flat_structure(dataset)))
# The incarnation ID is used to ensure consistency between the per-device
# iterators and the multi-device iterator.
@@ -216,13 +189,14 @@
for i, device in enumerate(self._devices):
ds = _PerDeviceGenerator(
i, self._multi_device_iterator_resource, self._incarnation_id,
- self._source_device_tensor, device, self._dataset.output_shapes,
- self._dataset.output_types, self._dataset.output_classes)
+ self._source_device_tensor, device, dataset._element_structure) # pylint: disable=protected-access
if prefetch_buffer_size > 0:
ds = ds.prefetch(prefetch_buffer_size)
- # TODO(jsimsa): Enable auto-tuning when supported for non-CPU devices.
+ # TODO(jsimsa): Enable auto-tuning and optimizations when supported for
+ # non-CPU devices.
options = dataset_ops.Options()
options.experimental_autotune = False
+ options.experimental_optimization.apply_default_optimizations = False
ds = ds.with_options(options)
with ops.device(device):
self._device_iterators.append(ds.make_initializable_iterator())
diff --git a/tensorflow/python/data/ops/optional_ops.py b/tensorflow/python/data/ops/optional_ops.py
index 15ec755..dcb743b 100644
--- a/tensorflow/python/data/ops/optional_ops.py
+++ b/tensorflow/python/data/ops/optional_ops.py
@@ -26,6 +26,7 @@
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
@six.add_metaclass(abc.ABCMeta)
@@ -145,6 +146,7 @@
return self._value_structure
+@tf_export("data.experimental.OptionalStructure")
class OptionalStructure(structure.Structure):
"""Represents an optional potentially containing a structured value."""
@@ -167,6 +169,10 @@
def _to_tensor_list(self, value):
return [value._variant_tensor] # pylint: disable=protected-access
+ def _to_batched_tensor_list(self, value):
+ raise NotImplementedError(
+ "Unbatching for `tf.data.experimental.Optional` objects.")
+
def _from_tensor_list(self, flat_value):
if (len(flat_value) != 1 or flat_value[0].dtype != dtypes.variant or
not flat_value[0].shape.is_compatible_with(tensor_shape.scalar())):
@@ -195,6 +201,10 @@
raise NotImplementedError(
"Batching for `tf.data.experimental.Optional` objects.")
+ def _unbatch(self):
+ raise NotImplementedError(
+ "Unbatching for `tf.data.experimental.Optional` objects.")
+
# pylint: disable=protected-access
structure.Structure._register_custom_converter(Optional,
diff --git a/tensorflow/python/data/ops/readers.py b/tensorflow/python/data/ops/readers.py
index 70a3b1b..0d6023d 100644
--- a/tensorflow/python/data/ops/readers.py
+++ b/tensorflow/python/data/ops/readers.py
@@ -20,6 +20,7 @@
from tensorflow.python.compat import compat
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.util import convert
+from tensorflow.python.data.util import structure
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
@@ -64,16 +65,8 @@
self._filenames, self._compression_type, self._buffer_size)
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.string
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
@tf_export(v1=["data.TextLineDataset"])
@@ -126,16 +119,8 @@
self._filenames, self._compression_type, self._buffer_size)
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.TensorShape([])
-
- @property
- def output_types(self):
- return dtypes.string
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
class ParallelInterleaveDataset(dataset_ops.InterleaveDataset):
@@ -168,7 +153,8 @@
self._buffer_output_elements,
self._prefetch_input_elements,
f=self._map_func.function,
- **dataset_ops.flat_structure(structure=self._output_structure))
+ **dataset_ops.flat_structure(self))
+ # pylint: enable=protected-access
def _transformation_name(self):
return "tf.data.experimental.parallel_interleave()"
@@ -247,16 +233,8 @@
return self._impl._inputs() # pylint: disable=protected-access
@property
- def output_classes(self):
- return self._impl.output_classes
-
- @property
- def output_shapes(self):
- return self._impl.output_shapes
-
- @property
- def output_types(self):
- return self._impl.output_types
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
@tf_export(v1=["data.TFRecordDataset"])
@@ -347,16 +325,8 @@
self._footer_bytes, self._buffer_size)
@property
- def output_classes(self):
- return ops.Tensor
-
- @property
- def output_shapes(self):
- return tensor_shape.scalar()
-
- @property
- def output_types(self):
- return dtypes.string
+ def _element_structure(self):
+ return structure.TensorStructure(dtypes.string, [])
@tf_export(v1=["data.FixedLengthRecordDataset"])
diff --git a/tensorflow/python/data/util/BUILD b/tensorflow/python/data/util/BUILD
index f15ebc3..04e8029 100644
--- a/tensorflow/python/data/util/BUILD
+++ b/tensorflow/python/data/util/BUILD
@@ -93,6 +93,7 @@
"//tensorflow/python:sparse_tensor",
"//tensorflow/python:tensor_shape",
"//tensorflow/python:variables",
+ "//tensorflow/python/data/kernel_tests:test_base",
"@absl_py//absl/testing:parameterized",
],
)
diff --git a/tensorflow/python/data/util/options.py b/tensorflow/python/data/util/options.py
index 9badba8..c578a83 100644
--- a/tensorflow/python/data/util/options.py
+++ b/tensorflow/python/data/util/options.py
@@ -48,27 +48,32 @@
return NotImplemented
-def create_option(name, ty, docstring, default=None):
+def create_option(name, ty, docstring, default_factory=lambda: None):
"""Creates a type-checked property.
Args:
- name: the name to use
- ty: the type to use
- docstring: the docstring to use
- default: the default value to use
+ name: The name to use.
+ ty: The type to use. The type of the property will be validated when it
+ is set.
+ docstring: The docstring to use.
+ default_factory: A callable that takes no arguments and returns a default
+ value to use if not set.
Returns:
A type-checked property.
"""
- def get_fn(self):
- return self._options.get(name, default) # pylint: disable=protected-access
+ def get_fn(option):
+ # pylint: disable=protected-access
+ if name not in option._options:
+ option._options[name] = default_factory()
+ return option._options.get(name)
- def set_fn(self, value):
+ def set_fn(option, value):
if not isinstance(value, ty):
raise TypeError("Property \"%s\" must be of type %s, got: %r (type: %r)" %
(name, ty, value, type(value)))
- self._options[name] = value # pylint: disable=protected-access
+ option._options[name] = value # pylint: disable=protected-access
return property(get_fn, set_fn, None, docstring)
diff --git a/tensorflow/python/data/util/options_test.py b/tensorflow/python/data/util/options_test.py
index c516983..8d41ef2 100644
--- a/tensorflow/python/data/util/options_test.py
+++ b/tensorflow/python/data/util/options_test.py
@@ -24,9 +24,12 @@
class _TestOptions(options.OptionsBase):
x = options.create_option(
- name="x", ty=int, docstring="the answer to everything", default=42)
+ name="x",
+ ty=int,
+ docstring="the answer to everything",
+ default_factory=lambda: 42)
y = options.create_option(
- name="y", ty=float, docstring="a tasty pie", default=3.14)
+ name="y", ty=float, docstring="a tasty pie", default_factory=lambda: 3.14)
class _NestedTestOptions(options.OptionsBase):
diff --git a/tensorflow/python/data/util/structure.py b/tensorflow/python/data/util/structure.py
index 874e8ab..9de0c4d 100644
--- a/tensorflow/python/data/util/structure.py
+++ b/tensorflow/python/data/util/structure.py
@@ -28,11 +28,13 @@
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import sparse_ops
+from tensorflow.python.util.tf_export import tf_export
_STRUCTURE_CONVERSION_FUNCTION_REGISTRY = {}
+@tf_export("data.experimental.Structure")
@six.add_metaclass(abc.ABCMeta)
class Structure(object):
"""Represents structural information, such as type and shape, about a value.
@@ -112,6 +114,26 @@
raise NotImplementedError("Structure._to_tensor_list()")
@abc.abstractmethod
+ def _to_batched_tensor_list(self, value):
+ """Returns a flat list of rank >= 1 `tf.Tensor` representing `value`.
+
+ This method can be used, along with `self._flat_shapes` and
+ `self._flat_types` to represent structured values in lower level APIs
+ (such as plain TensorFlow operations) that do not understand structure,
+ *and* that require that the plain tensors have a rank of at least one
+ (e.g. for the purpose of slicing the tensors).
+
+ Requires: `self.is_compatible_with(Structure.from_value(value))`.
+
+ Args:
+ value: A value with compatible structure.
+
+ Returns:
+ A flat list of `tf.Tensor` representing `value`.
+ """
+ raise NotImplementedError("Structure._to_batched_tensor_list()")
+
+ @abc.abstractmethod
def _from_tensor_list(self, flat_value):
"""Builds a flat list of `tf.Tensor` into a value matching this structure.
@@ -157,6 +179,10 @@
"""
raise NotImplementedError("Structure._batch()")
+ @abc.abstractmethod
+ def _unbatch(self):
+ raise NotImplementedError("Structure._unbatch()")
+
@staticmethod
def from_value(value):
"""Returns a `Structure` that represents the given `value`.
@@ -191,56 +217,6 @@
return TensorStructure.from_value(tensor)
@staticmethod
- def _from_legacy_structure(output_types, output_shapes, output_classes):
- """Returns a `Structure` that represents the given legacy structure.
-
- This method provides a way to convert from the existing `Dataset` and
- `Iterator` structure-related properties to a `Structure` object.
-
- TODO(b/110122868): Remove this method once `Structure` is used throughout
- `tf.data`.
-
- Args:
- output_types: A nested structure of `tf.DType` objects corresponding to
- each component of a structured value.
- output_shapes: A nested structure of `tf.TensorShape` objects
- corresponding to each component a structured value.
- output_classes: A nested structure of Python `type` objects corresponding
- to each component of a structured value.
-
- Returns:
- A `Structure`.
-
- Raises:
- TypeError: If a structure cannot be built the arguments, because one of
- the component classes in `output_classes` is not supported.
- """
- flat_types = nest.flatten(output_types)
- flat_shapes = nest.flatten(output_shapes)
- flat_classes = nest.flatten(output_classes)
- flat_ret = []
- for flat_type, flat_shape, flat_class in zip(flat_types, flat_shapes,
- flat_classes):
- if isinstance(flat_class, Structure):
- flat_ret.append(flat_class)
- elif issubclass(flat_class, sparse_tensor_lib.SparseTensor):
- flat_ret.append(SparseTensorStructure(flat_type, flat_shape))
- elif issubclass(flat_class, ops.Tensor):
- flat_ret.append(TensorStructure(flat_type, flat_shape))
- else:
- # NOTE(mrry): Since legacy structures produced by iterators only
- # comprise Tensors, SparseTensors, and nests, we do not need to
- # support all structure types here.
- raise TypeError(
- "Could not build a structure for output class %r" % flat_type)
-
- ret = nest.pack_sequence_as(output_classes, flat_ret)
- if isinstance(ret, Structure):
- return ret
- else:
- return NestedStructure(ret)
-
- @staticmethod
def _register_custom_converter(type_object, converter_fn):
"""Registers `converter_fn` for converting values of the given type.
@@ -265,9 +241,63 @@
raise NotImplementedError("Structure._to_legacy_output_classes()")
+def convert_legacy_structure(output_types, output_shapes, output_classes):
+ """Returns a `Structure` that represents the given legacy structure.
+
+ This method provides a way to convert from the existing `Dataset` and
+ `Iterator` structure-related properties to a `Structure` object. A "legacy"
+ structure is represented by the `tf.data.Dataset.output_types`,
+ `tf.data.Dataset.output_shapes`, and `tf.data.Dataset.output_classes`
+ properties.
+
+ TODO(b/110122868): Remove this function once `Structure` is used throughout
+ `tf.data`.
+
+ Args:
+ output_types: A nested structure of `tf.DType` objects corresponding to
+ each component of a structured value.
+ output_shapes: A nested structure of `tf.TensorShape` objects
+ corresponding to each component a structured value.
+ output_classes: A nested structure of Python `type` objects corresponding
+ to each component of a structured value.
+
+ Returns:
+ A `Structure`.
+
+ Raises:
+ TypeError: If a structure cannot be built from the arguments, because one of
+ the component classes in `output_classes` is not supported.
+ """
+ flat_types = nest.flatten(output_types)
+ flat_shapes = nest.flatten(output_shapes)
+ flat_classes = nest.flatten(output_classes)
+ flat_ret = []
+ for flat_type, flat_shape, flat_class in zip(flat_types, flat_shapes,
+ flat_classes):
+ if isinstance(flat_class, Structure):
+ flat_ret.append(flat_class)
+ elif issubclass(flat_class, sparse_tensor_lib.SparseTensor):
+ flat_ret.append(SparseTensorStructure(flat_type, flat_shape))
+ elif issubclass(flat_class, ops.Tensor):
+ flat_ret.append(TensorStructure(flat_type, flat_shape))
+ else:
+ # NOTE(mrry): Since legacy structures produced by iterators only
+ # comprise Tensors, SparseTensors, and nests, we do not need to
+ # support all structure types here.
+ raise TypeError(
+ "Could not build a structure for output class %r" % flat_type)
+
+ ret = nest.pack_sequence_as(output_classes, flat_ret)
+ if isinstance(ret, Structure):
+ return ret
+ else:
+ return NestedStructure(ret)
+
+
# NOTE(mrry): The following classes make extensive use of non-public methods of
# their base class, so we disable the protected-access lint warning once here.
# pylint: disable=protected-access
+@tf_export("data.experimental.NestedStructure")
class NestedStructure(Structure):
"""Represents a nested structure in which each leaf is a `Structure`."""
@@ -323,6 +353,22 @@
ret.extend(structure._to_tensor_list(sub_value))
return ret
+ def _to_batched_tensor_list(self, value):
+ ret = []
+
+ try:
+ flat_value = nest.flatten_up_to(self._nested_structure, value)
+ except (ValueError, TypeError):
+ raise ValueError("The value %r is not compatible with the nested "
+ "structure %r." % (value, self._nested_structure))
+
+ for sub_value, structure in zip(flat_value, self._flat_nested_structure):
+ if not structure.is_compatible_with(Structure.from_value(sub_value)):
+ raise ValueError("Component value %r is not compatible with the nested "
+ "structure %r." % (sub_value, structure))
+ ret.extend(structure._to_batched_tensor_list(sub_value))
+ return ret
+
def _from_tensor_list(self, flat_value):
if len(flat_value) != len(self._flat_types):
raise ValueError("Expected %d flat values in NestedStructure but got %d."
@@ -372,7 +418,12 @@
return NestedStructure(nest.map_structure(
lambda s: s._batch(batch_size), self._nested_structure))
+ def _unbatch(self):
+ return NestedStructure(nest.map_structure(
+ lambda s: s._unbatch(), self._nested_structure))
+
+@tf_export("data.experimental.TensorStructure")
class TensorStructure(Structure):
"""Represents structural information about a `tf.Tensor`."""
@@ -399,6 +450,11 @@
"and shape %s." % (value, self._dtype, self._shape))
return [value]
+ def _to_batched_tensor_list(self, value):
+ if self._shape.merge_with(value.shape).ndims == 0:
+ raise ValueError("Unbatching a tensor is only supported for rank >= 1")
+ return [value]
+
def _from_tensor_list(self, flat_value):
if len(flat_value) != 1:
raise ValueError("TensorStructure corresponds to a single tf.Tensor.")
@@ -435,7 +491,13 @@
self._dtype,
tensor_shape.TensorShape([batch_size]).concatenate(self._shape))
+ def _unbatch(self):
+ if self._shape.ndims == 0:
+ raise ValueError("Unbatching a tensor is only supported for rank >= 1")
+ return TensorStructure(self._dtype, self._shape[1:])
+
+@tf_export("data.experimental.SparseTensorStructure")
class SparseTensorStructure(Structure):
"""Represents structural information about a `tf.SparseTensor`."""
@@ -463,6 +525,13 @@
def _to_tensor_list(self, value):
return [sparse_ops.serialize_sparse(value, out_type=dtypes.variant)]
+ def _to_batched_tensor_list(self, value):
+ if self._dense_shape.merge_with(
+ tensor_util.constant_value_as_shape(value.dense_shape)).ndims == 0:
+ raise ValueError(
+ "Unbatching a sparse tensor is only supported for rank >= 1")
+ return [sparse_ops.serialize_many_sparse(value, out_type=dtypes.variant)]
+
def _from_tensor_list(self, flat_value):
if (len(flat_value) != 1 or flat_value[0].dtype != dtypes.variant or
not flat_value[0].shape.is_compatible_with(tensor_shape.vector(3))):
@@ -497,3 +566,8 @@
return SparseTensorStructure(
self._dtype,
tensor_shape.TensorShape([batch_size]).concatenate(self._dense_shape))
+
+ def _unbatch(self):
+ if self._dense_shape.ndims == 0:
+ raise ValueError("Unbatching a tensor is only supported for rank >= 1")
+ return SparseTensorStructure(self._dtype, self._dense_shape[1:])
diff --git a/tensorflow/python/data/util/structure_test.py b/tensorflow/python/data/util/structure_test.py
index 314a459..91dcfa6 100644
--- a/tensorflow/python/data/util/structure_test.py
+++ b/tensorflow/python/data/util/structure_test.py
@@ -21,6 +21,7 @@
from absl.testing import parameterized
import numpy as np
+from tensorflow.python.data.kernel_tests import test_base
from tensorflow.python.data.util import nest
from tensorflow.python.data.util import structure
from tensorflow.python.framework import constant_op
@@ -34,7 +35,7 @@
from tensorflow.python.platform import test
-class StructureTest(test.TestCase, parameterized.TestCase):
+class StructureTest(test_base.DatasetTestBase, parameterized.TestCase):
# NOTE(mrry): The arguments must be lifted into lambdas because otherwise they
# will be executed before the (eager- or graph-mode) test environment has been
@@ -352,9 +353,9 @@
"b": (structure.SparseTensorStructure(dtypes.int32, [2, 2]),
structure.TensorStructure(dtypes.string, []))})),
)
- def testFromLegacyStructure(self, output_types, output_shapes, output_classes,
- expected_structure):
- actual_structure = structure.Structure._from_legacy_structure(
+ def testConvertLegacyStructure(self, output_types, output_shapes,
+ output_classes, expected_structure):
+ actual_structure = structure.convert_legacy_structure(
output_types, output_shapes, output_classes)
self.assertTrue(expected_structure.is_compatible_with(actual_structure))
self.assertTrue(actual_structure.is_compatible_with(expected_structure))
@@ -418,6 +419,75 @@
self.assertTrue(
expected_batched_structure.is_compatible_with(batched_structure))
+ @parameterized.named_parameters(
+ ("Tensor", structure.TensorStructure(dtypes.float32, [32]),
+ structure.TensorStructure(dtypes.float32, [])),
+ ("TensorUnknown", structure.TensorStructure(dtypes.float32, [None]),
+ structure.TensorStructure(dtypes.float32, [])),
+ ("SparseTensor",
+ structure.SparseTensorStructure(dtypes.float32, [32, None]),
+ structure.SparseTensorStructure(dtypes.float32, [None])),
+ ("SparseTensorUnknown",
+ structure.SparseTensorStructure(dtypes.float32, [None, 4]),
+ structure.SparseTensorStructure(dtypes.float32, [4])),
+ ("Nest", structure.NestedStructure({
+ "a": structure.TensorStructure(dtypes.float32, [128]),
+ "b": (structure.SparseTensorStructure(dtypes.int32, [128, 2, 2]),
+ structure.TensorStructure(dtypes.string, [None]))}),
+ structure.NestedStructure({
+ "a": structure.TensorStructure(dtypes.float32, []),
+ "b": (structure.SparseTensorStructure(dtypes.int32, [2, 2]),
+ structure.TensorStructure(dtypes.string, []))})),
+ )
+ def testUnbatch(self, element_structure, expected_unbatched_structure):
+ unbatched_structure = element_structure._unbatch()
+ self.assertTrue(
+ unbatched_structure.is_compatible_with(expected_unbatched_structure))
+ self.assertTrue(
+ expected_unbatched_structure.is_compatible_with(unbatched_structure))
+
+ # pylint: disable=g-long-lambda
+ @parameterized.named_parameters(
+ ("Tensor", lambda: constant_op.constant([[1.0, 2.0], [3.0, 4.0]]),
+ lambda: constant_op.constant([1.0, 2.0])),
+ ("SparseTensor", lambda: sparse_tensor.SparseTensor(
+ indices=[[0, 0], [1, 1]], values=[13, 27], dense_shape=[2, 2]),
+ lambda: sparse_tensor.SparseTensor(
+ indices=[[0]], values=[13], dense_shape=[2])),
+ ("Nest", lambda: (
+ constant_op.constant([[1.0, 2.0], [3.0, 4.0]]),
+ sparse_tensor.SparseTensor(
+ indices=[[0, 0], [1, 1]], values=[13, 27], dense_shape=[2, 2])),
+ lambda: (constant_op.constant([1.0, 2.0]), sparse_tensor.SparseTensor(
+ indices=[[0]], values=[13], dense_shape=[2]))),
+ )
+ def testToBatchedTensorList(self, value_fn, element_0_fn):
+ batched_value = value_fn()
+ s = structure.Structure.from_value(batched_value)
+ batched_tensor_list = s._to_batched_tensor_list(batched_value)
+
+ # The batch dimension is 2 for all of the test cases.
+ # NOTE(mrry): `tf.shape()` does not currently work for the DT_VARIANT
+ # tensors in which we store sparse tensors.
+ for t in batched_tensor_list:
+ if t.dtype != dtypes.variant:
+ self.assertEqual(2, self.evaluate(array_ops.shape(t)[0]))
+
+ # Test that the 0th element from the unbatched tensor is equal to the
+ # expected value.
+ expected_element_0 = self.evaluate(element_0_fn())
+ unbatched_s = s._unbatch()
+ actual_element_0 = unbatched_s._from_tensor_list(
+ [t[0] for t in batched_tensor_list])
+
+ for expected, actual in zip(
+ nest.flatten(expected_element_0), nest.flatten(actual_element_0)):
+ if sparse_tensor.is_sparse(expected):
+ self.assertSparseValuesEqual(expected, actual)
+ else:
+ self.assertAllEqual(expected, actual)
+
+ # pylint: enable=g-long-lambda
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/debug/BUILD b/tensorflow/python/debug/BUILD
index c6abd47..1dcdb88 100644
--- a/tensorflow/python/debug/BUILD
+++ b/tensorflow/python/debug/BUILD
@@ -1132,4 +1132,7 @@
":debug_tflearn_iris",
":offline_analyzer",
],
+ tags = [
+ "no_windows",
+ ],
)
diff --git a/tensorflow/python/debug/cli/analyzer_cli_test.py b/tensorflow/python/debug/cli/analyzer_cli_test.py
index 322ecf9..586982d 100644
--- a/tensorflow/python/debug/cli/analyzer_cli_test.py
+++ b/tensorflow/python/debug/cli/analyzer_cli_test.py
@@ -573,6 +573,7 @@
return analyzer, registry
+@test_util.run_v1_only("b/120545219")
class AnalyzerCLISimpleMulAddTest(test_util.TensorFlowTestCase):
@classmethod
@@ -645,7 +646,6 @@
self.assertEqual(len("Size (B)") + 1, dump_size_col_width)
self.assertEqual(len("Op type") + 1, op_type_col_width)
- @test_util.run_deprecated_v1
def testMeasureTensorListColumnWidthsGivesRightAnswerForData(self):
dump = self._debug_dump.dumped_tensor_data[0]
self.assertLess(dump.dump_size_bytes, 1000)
@@ -661,7 +661,6 @@
# column should be determined by the length of "VariableV2".
self.assertEqual(len("VariableV2") + 1, op_type_col_width)
- @test_util.run_deprecated_v1
def testListTensors(self):
# Use shorthand alias for the command prefix.
out = self._registry.dispatch_command("lt", [])
@@ -675,7 +674,6 @@
# Check the main menu.
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsInReverseTimeOrderWorks(self):
# Use shorthand alias for the command prefix.
out = self._registry.dispatch_command("lt", ["-s", "timestamp", "-r"])
@@ -691,7 +689,6 @@
reverse=True)
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsInDumpSizeOrderWorks(self):
out = self._registry.dispatch_command("lt", ["-s", "dump_size"])
assert_listed_tensors(
@@ -705,7 +702,6 @@
sort_by="dump_size")
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsInReverseDumpSizeOrderWorks(self):
out = self._registry.dispatch_command("lt", ["-s", "dump_size", "-r"])
assert_listed_tensors(
@@ -725,7 +721,6 @@
self.assertIn("ValueError: Unsupported key to sort tensors by: foobar",
out.lines)
- @test_util.run_deprecated_v1
def testListTensorsInOpTypeOrderWorks(self):
# Use shorthand alias for the command prefix.
out = self._registry.dispatch_command("lt", ["-s", "op_type"])
@@ -741,7 +736,6 @@
reverse=False)
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsInReverseOpTypeOrderWorks(self):
# Use shorthand alias for the command prefix.
out = self._registry.dispatch_command("lt", ["-s", "op_type", "-r"])
@@ -757,7 +751,6 @@
reverse=True)
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsInTensorNameOrderWorks(self):
# Use shorthand alias for the command prefix.
out = self._registry.dispatch_command("lt", ["-s", "tensor_name"])
@@ -773,7 +766,6 @@
reverse=False)
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsInReverseTensorNameOrderWorks(self):
# Use shorthand alias for the command prefix.
out = self._registry.dispatch_command("lt", ["-s", "tensor_name", "-r"])
@@ -789,7 +781,6 @@
reverse=True)
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorsFilterByNodeNameRegex(self):
out = self._registry.dispatch_command("list_tensors",
["--node_name_filter", ".*read.*"])
@@ -803,7 +794,6 @@
assert_listed_tensors(self, out, [], [], node_name_regex="^read")
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorFilterByOpTypeRegex(self):
out = self._registry.dispatch_command("list_tensors",
["--op_type_filter", "Identity"])
@@ -832,7 +822,6 @@
op_type_regex="(Add|MatMul)")
check_main_menu(self, out, list_tensors_enabled=False)
- @test_util.run_deprecated_v1
def testListTensorWithFilterAndNodeNameExclusionWorks(self):
# First, create and register the filter.
def is_2x1_vector(datum, tensor):
@@ -889,7 +878,6 @@
out = self._registry.dispatch_command("list_tensors", ["--bar"])
check_syntax_error_output(self, out, "list_tensors")
- @test_util.run_deprecated_v1
def testNodeInfoByNodeName(self):
node_name = "simple_mul_add/matmul"
out = self._registry.dispatch_command("node_info", [node_name])
@@ -914,7 +902,6 @@
[(len(out.lines[0]) - len(node_name), len(out.lines[0]), "bold")],
out.font_attr_segs[0])
- @test_util.run_deprecated_v1
def testNodeInfoShowAttributes(self):
node_name = "simple_mul_add/matmul"
out = self._registry.dispatch_command("node_info", ["-a", node_name])
@@ -938,7 +925,6 @@
print_tensor_node_name=node_name,
list_outputs_node_name=node_name)
- @test_util.run_deprecated_v1
def testNodeInfoShowDumps(self):
node_name = "simple_mul_add/matmul"
out = self._registry.dispatch_command("node_info", ["-d", node_name])
@@ -963,7 +949,6 @@
len(out.lines[16]) - len(out.lines[16].strip()),
len(out.lines[16]), "pt %s:0 -n 0" % node_name)
- @test_util.run_deprecated_v1
def testNodeInfoShowStackTraceUnavailableIsIndicated(self):
self._debug_dump.set_python_graph(None)
@@ -987,7 +972,6 @@
print_tensor_node_name=node_name,
list_outputs_node_name=node_name)
- @test_util.run_deprecated_v1
def testNodeInfoShowStackTraceAvailableWorks(self):
self._debug_dump.set_python_graph(self._sess.graph)
@@ -1011,7 +995,6 @@
print_tensor_node_name=node_name,
list_outputs_node_name=node_name)
- @test_util.run_deprecated_v1
def testNodeInfoByTensorName(self):
node_name = "simple_mul_add/u/read"
tensor_name = node_name + ":0"
@@ -1381,7 +1364,6 @@
break
return index
- @test_util.run_deprecated_v1
def testPrintSourceForOpNamesWholeFileWorks(self):
self._debug_dump.set_python_graph(self._sess.graph)
out = self._registry.dispatch_command(
@@ -1434,7 +1416,6 @@
self.assertEqual("pt simple_mul_add/add",
out.font_attr_segs[index + 1][0][2].content)
- @test_util.run_deprecated_v1
def testPrintSourceForTensorNamesWholeFileWorks(self):
self._debug_dump.set_python_graph(self._sess.graph)
out = self._registry.dispatch_command(
@@ -1455,7 +1436,6 @@
self.assertEqual("pt simple_mul_add/u:0",
out.font_attr_segs[index + 2][0][2].content)
- @test_util.run_deprecated_v1
def testPrintSourceForOpNamesStartingAtSpecifiedLineWorks(self):
self._debug_dump.set_python_graph(self._sess.graph)
out = self._registry.dispatch_command(
@@ -1482,7 +1462,6 @@
self.assertEqual("pt simple_mul_add/u/read",
out.font_attr_segs[index + 3][0][2].content)
- @test_util.run_deprecated_v1
def testPrintSourceForOpNameSettingMaximumElementCountWorks(self):
self._debug_dump.set_python_graph(self._sess.graph)
out = self._registry.dispatch_command(
@@ -1527,7 +1506,6 @@
self.assertTrue(cli_shared.COLOR_GRAY in attr_seg[2] or
attr_seg[2] == cli_shared.COLOR_GRAY)
- @test_util.run_deprecated_v1
def testListSourceWithNodeNameFilterWithMatchesWorks(self):
self._debug_dump.set_python_graph(self._sess.graph)
out = self._registry.dispatch_command("list_source", ["-n", ".*/read"])
@@ -1691,6 +1669,7 @@
self.assertNotIn("...,", out.lines[4])
+@test_util.run_v1_only("b/120545219")
class AnalyzerCLIControlDepTest(test_util.TensorFlowTestCase):
@classmethod
@@ -1742,7 +1721,6 @@
# Tear down temporary dump directory.
shutil.rmtree(cls._dump_root)
- @test_util.run_deprecated_v1
def testNodeInfoWithControlDependencies(self):
# Call node_info on a node with control inputs.
out = self._registry.dispatch_command("node_info",
@@ -1783,7 +1761,6 @@
len(out.lines[z_line]),
"ni -a -d -t control_deps/ctrl_dep_z")
- @test_util.run_deprecated_v1
def testListInputsNonRecursiveNoControl(self):
"""List inputs non-recursively, without any control inputs."""
@@ -1826,7 +1803,6 @@
len(out.lines[3]) - len("control_deps/ctrl_dep_y"),
len(out.lines[3]), "li -c -r control_deps/ctrl_dep_y")
- @test_util.run_deprecated_v1
def testListInputsNonRecursiveNoControlUsingTensorName(self):
"""List inputs using the name of an output tensor of the node."""
@@ -1855,7 +1831,6 @@
len(out.lines[3]) - len("control_deps/ctrl_dep_y"),
len(out.lines[3]), "li -c -r control_deps/ctrl_dep_y")
- @test_util.run_deprecated_v1
def testListInputsNonRecursiveWithControls(self):
"""List inputs non-recursively, with control inputs."""
node_name = "control_deps/ctrl_dep_z"
@@ -1886,7 +1861,6 @@
len(out.lines[5]) - len("control_deps/x"),
len(out.lines[5]), "li -c -r control_deps/x")
- @test_util.run_deprecated_v1
def testListInputsRecursiveWithControls(self):
"""List inputs recursively, with control inputs."""
node_name = "control_deps/ctrl_dep_z"
@@ -1932,7 +1906,6 @@
len(out.lines[18]) - len("control_deps/x"),
len(out.lines[18]), "li -c -r control_deps/x")
- @test_util.run_deprecated_v1
def testListInputsRecursiveWithControlsWithDepthLimit(self):
"""List inputs recursively, with control inputs and a depth limit."""
node_name = "control_deps/ctrl_dep_z"
@@ -1992,7 +1965,6 @@
"ERROR: There is no node named \"control_deps/z/foo\" in the "
"partition graphs"], out.lines)
- @test_util.run_deprecated_v1
def testListRecipientsRecursiveWithControlsWithDepthLimit(self):
"""List recipients recursively, with control inputs and a depth limit."""
@@ -2025,6 +1997,7 @@
out.font_attr_segs[0])
+@test_util.run_v1_only("b/120545219")
class AnalyzerCLIWhileLoopTest(test_util.TensorFlowTestCase):
@classmethod
@@ -2064,7 +2037,6 @@
# Tear down temporary dump directory.
shutil.rmtree(cls._dump_root)
- @test_util.run_deprecated_v1
def testMultipleDumpsPrintTensorNoNumber(self):
output = self._registry.dispatch_command("pt", ["while/Identity:0"])
@@ -2082,7 +2054,6 @@
self.assertEqual("For example:", output.lines[-2])
self.assertEqual(" print_tensor while/Identity:0 -n 0", output.lines[-1])
- @test_util.run_deprecated_v1
def testMultipleDumpsPrintTensorWithNumber(self):
for i in xrange(5):
output = self._registry.dispatch_command(
@@ -2096,7 +2067,6 @@
self.assertTrue(output.lines[4].startswith("array(%d" % i))
self.assertTrue(output.lines[4].endswith(")"))
- @test_util.run_deprecated_v1
def testMultipleDumpsPrintTensorInvalidNumber(self):
output = self._registry.dispatch_command("pt",
["while/Identity:0", "-n", "10"])
diff --git a/tensorflow/python/debug/cli/cli_shared_test.py b/tensorflow/python/debug/cli/cli_shared_test.py
index d191a23..66a12ef 100644
--- a/tensorflow/python/debug/cli/cli_shared_test.py
+++ b/tensorflow/python/debug/cli/cli_shared_test.py
@@ -105,6 +105,7 @@
cli_shared.time_to_readable_str(100, force_time_unit="ks")
+@test_util.run_v1_only("b/120545219")
class GetRunStartIntroAndDescriptionTest(test_util.TensorFlowTestCase):
def setUp(self):
@@ -334,6 +335,7 @@
self.assertEqual("run #1: 1 fetch (a:0); 1 feed (foo)", short_description)
+@test_util.run_v1_only("b/120545219")
class GetErrorIntroTest(test_util.TensorFlowTestCase):
def setUp(self):
@@ -342,7 +344,6 @@
def tearDown(self):
ops.reset_default_graph()
- @test_util.run_deprecated_v1
def testShapeError(self):
tf_error = errors.OpError(None, self.var_a.initializer, "foo description",
None)
diff --git a/tensorflow/python/debug/cli/profile_analyzer_cli_test.py b/tensorflow/python/debug/cli/profile_analyzer_cli_test.py
index effcd50..d6d2b58 100644
--- a/tensorflow/python/debug/cli/profile_analyzer_cli_test.py
+++ b/tensorflow/python/debug/cli/profile_analyzer_cli_test.py
@@ -70,6 +70,7 @@
"%s matched at least one line in %s." % (pattern, str(lines)))
+@test_util.run_v1_only("b/120545219")
class ProfileAnalyzerListProfileTest(test_util.TensorFlowTestCase):
def testNodeInfoEmpty(self):
@@ -321,6 +322,7 @@
_assert_at_least_one_line_matches(r"Device Total.*0\.009ms", prof_output)
+@test_util.run_v1_only("b/120545219")
class ProfileAnalyzerPrintSourceTest(test_util.TensorFlowTestCase):
def setUp(self):
@@ -348,7 +350,6 @@
ops.reset_default_graph()
super(ProfileAnalyzerPrintSourceTest, self).tearDown()
- @test_util.run_deprecated_v1
def testPrintSourceForWhileLoop(self):
prof_output = self.prof_analyzer.print_source([__file__])
@@ -362,7 +363,6 @@
r"\[(\|)+(\s)*\] .*us .*7\(55\) .*L%d.*(\S)+" % self.loop_lineno,
prof_output.lines)
- @test_util.run_deprecated_v1
def testPrintSourceOutputContainsClickableLinks(self):
prof_output = self.prof_analyzer.print_source([__file__])
any_match, line_index = _at_least_one_line_matches(
@@ -379,7 +379,6 @@
break
self.assertTrue(any_menu_item_match)
- @test_util.run_deprecated_v1
def testPrintSourceWithNonDefaultTimeUnit(self):
prof_output = self.prof_analyzer.print_source([
__file__, "--time_unit", "ms"])
@@ -394,7 +393,6 @@
r"\[(\|)+(\s)*\] .*ms .*7\(55\) .*L%d.*(\S)+" % self.loop_lineno,
prof_output.lines)
- @test_util.run_deprecated_v1
def testPrintSourceWithNodeNameFilter(self):
prof_output = self.prof_analyzer.print_source([
__file__, "--node_name_filter", "x$"])
@@ -427,7 +425,6 @@
break
self.assertTrue(any_menu_item_match)
- @test_util.run_deprecated_v1
def testPrintSourceWithOpTypeFilter(self):
prof_output = self.prof_analyzer.print_source([
__file__, "--op_type_filter", "Less"])
diff --git a/tensorflow/python/debug/cli/stepper_cli_test.py b/tensorflow/python/debug/cli/stepper_cli_test.py
index 7b8a42c..5cf69d0 100644
--- a/tensorflow/python/debug/cli/stepper_cli_test.py
+++ b/tensorflow/python/debug/cli/stepper_cli_test.py
@@ -129,6 +129,7 @@
return updated
+@test_util.run_v1_only("b/120545219")
class NodeStepperSimpleGraphTest(test_util.TensorFlowTestCase):
def setUp(self):
diff --git a/tensorflow/python/debug/examples/debug_errors.py b/tensorflow/python/debug/examples/debug_errors.py
index 28abc97..e369207 100644
--- a/tensorflow/python/debug/examples/debug_errors.py
+++ b/tensorflow/python/debug/examples/debug_errors.py
@@ -77,4 +77,5 @@
default=False,
help="Use debugger to track down bad values during training")
FLAGS, unparsed = parser.parse_known_args()
- tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
+ with tf.Graph().as_default():
+ tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
diff --git a/tensorflow/python/debug/examples/debug_fibonacci.py b/tensorflow/python/debug/examples/debug_fibonacci.py
index 3821b39..777fb08 100644
--- a/tensorflow/python/debug/examples/debug_fibonacci.py
+++ b/tensorflow/python/debug/examples/debug_fibonacci.py
@@ -100,4 +100,5 @@
"--debug flag.")
FLAGS, unparsed = parser.parse_known_args()
- tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
+ with tf.Graph().as_default():
+ tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
diff --git a/tensorflow/python/debug/examples/debug_keras.py b/tensorflow/python/debug/examples/debug_keras.py
index 3272d85..019121f 100644
--- a/tensorflow/python/debug/examples/debug_keras.py
+++ b/tensorflow/python/debug/examples/debug_keras.py
@@ -86,4 +86,5 @@
default=2,
help="Number of epochs to train the model for.")
FLAGS, unparsed = parser.parse_known_args()
- tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
+ with tf.Graph().as_default():
+ tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
diff --git a/tensorflow/python/debug/examples/debug_mnist.py b/tensorflow/python/debug/examples/debug_mnist.py
index ab1c903..09fb06c 100644
--- a/tensorflow/python/debug/examples/debug_mnist.py
+++ b/tensorflow/python/debug/examples/debug_mnist.py
@@ -190,4 +190,5 @@
"the gRPC address (e.g., localhost:1234). Mutually exclusive with the "
"--debug flag.")
FLAGS, unparsed = parser.parse_known_args()
- tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
+ with tf.Graph().as_default():
+ tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
diff --git a/tensorflow/python/debug/lib/debug_gradients_test.py b/tensorflow/python/debug/lib/debug_gradients_test.py
index 1c53147..885691c 100644
--- a/tensorflow/python/debug/lib/debug_gradients_test.py
+++ b/tensorflow/python/debug/lib/debug_gradients_test.py
@@ -36,6 +36,7 @@
from tensorflow.python.training import gradient_descent
+@test_util.run_v1_only("b/120545219")
class IdentifyGradientTest(test_util.TensorFlowTestCase):
def setUp(self):
@@ -54,7 +55,6 @@
ops.reset_default_graph()
debug_gradients.clear_gradient_debuggers()
- @test_util.run_deprecated_v1
def testIdentifyGradientGivesCorrectTensorObjectWithoutContextManager(self):
grad_debugger = debug_gradients.GradientsDebugger()
id_grad_w = grad_debugger.identify_gradient(self.w)
@@ -85,7 +85,6 @@
self.assertIsInstance(w_grad, ops.Tensor)
self.assertAllClose(1.0, self.sess.run(w_grad))
- @test_util.run_deprecated_v1
def testIdentifyGradientGivesCorrectTensorObjectWithTfGradients(self):
grad_debugger = debug_gradients.GradientsDebugger()
id_grad_w = grad_debugger.identify_gradient(self.w)
@@ -117,7 +116,6 @@
self.assertIsInstance(w_grad, ops.Tensor)
self.assertAllClose(1.0, self.sess.run(w_grad))
- @test_util.run_deprecated_v1
def testCallingIdentifyGradientTwiceWithTheSameGradientsDebuggerErrors(self):
grad_debugger = debug_gradients.GradientsDebugger()
grad_debugger.identify_gradient(self.w)
@@ -125,7 +123,6 @@
"The graph already contains an op named .*"):
grad_debugger.identify_gradient(self.w)
- @test_util.run_deprecated_v1
def testIdentifyGradientWorksOnMultipleLosses(self):
grad_debugger_1 = debug_gradients.GradientsDebugger()
grad_debugger_2 = debug_gradients.GradientsDebugger()
@@ -154,7 +151,6 @@
self.assertAllClose(2.0 * 5.0, self.sess.run(dz1_dy))
self.assertAllClose(0.5 * (5.0**-0.5), self.sess.run(dz2_dy))
- @test_util.run_deprecated_v1
def testIdentifyGradientRaisesLookupErrorForUnknownXTensor(self):
grad_debugger_1 = debug_gradients.GradientsDebugger()
grad_debugger_2 = debug_gradients.GradientsDebugger()
@@ -175,7 +171,6 @@
r"This GradientsDebugger has not received any gradient tensor for "):
grad_debugger_2.gradient_tensor(self.w)
- @test_util.run_deprecated_v1
def testIdentifyGradientRaisesTypeErrorForNonTensorOrTensorNameInput(self):
grad_debugger = debug_gradients.GradientsDebugger()
with self.assertRaisesRegexp(
@@ -184,7 +179,6 @@
r"has type .*Operation.*"):
grad_debugger.gradient_tensor(variables.global_variables_initializer())
- @test_util.run_deprecated_v1
def testIdentifyGradientTensorWorksWithGradientDescentOptimizer(self):
grad_debugger = debug_gradients.GradientsDebugger()
id_grad_w = grad_debugger.identify_gradient(self.w)
@@ -200,7 +194,6 @@
self.assertIsInstance(w_grad, ops.Tensor)
self.assertAllClose(1.0, self.sess.run(w_grad))
- @test_util.run_deprecated_v1
def testWatchGradientsByXTensorNamesWorks(self):
y = math_ops.add(self.w, -1.0, name="y")
@@ -227,7 +220,6 @@
self.assertIsInstance(w_grad, ops.Tensor)
self.assertAllClose(1.0, self.sess.run(w_grad))
- @test_util.run_deprecated_v1
def testWatchGradientsByXTensorNamesWorksWithoutContextManager(self):
y = math_ops.add(self.w, -1.0, name="y")
@@ -254,7 +246,6 @@
self.assertIsInstance(w_grad, ops.Tensor)
self.assertAllClose(1.0, self.sess.run(w_grad))
- @test_util.run_deprecated_v1
def testWatchGradientsWorksOnRefTensor(self):
y = math_ops.add(self.w, -1.0, name="y")
@@ -273,7 +264,6 @@
self.assertAllClose(3.0, self.sess.run(
grad_debugger.gradient_tensor("u:0")))
- @test_util.run_deprecated_v1
def testWatchGradientsWorksOnMultipleTensors(self):
y = math_ops.add(self.w, -1.0, name="y")
@@ -294,7 +284,6 @@
self.assertAllClose(3.0, self.sess.run(
grad_debugger.gradient_tensor("u:0")))
- @test_util.run_deprecated_v1
def testWatchGradientsByXTensorsWorks(self):
y = math_ops.add(self.w, -1.0, name="foo/y")
z = math_ops.square(y, name="foo/z")
@@ -317,7 +306,6 @@
self.assertAllClose(10.0, self.sess.run(w_grad))
self.assertAllClose(30.0, self.sess.run(u_grad))
- @test_util.run_deprecated_v1
def testWatchGradientsByTensorCanWorkOnMultipleLosses(self):
y = math_ops.add(self.w, -1.0, name="y")
z1 = math_ops.square(y, name="z1")
@@ -343,7 +331,6 @@
self.assertAllClose(2.0 * 5.0, self.sess.run(dz1_dy))
self.assertAllClose(0.5 * (5.0**-0.5), self.sess.run(dz2_dy))
- @test_util.run_deprecated_v1
def testGradientsValuesFromDumpWorks(self):
y = math_ops.add(self.w, -1.0, name="y")
z = math_ops.square(y, name="z")
diff --git a/tensorflow/python/debug/lib/debug_utils_test.py b/tensorflow/python/debug/lib/debug_utils_test.py
index cf59b30..9d59cfc 100644
--- a/tensorflow/python/debug/lib/debug_utils_test.py
+++ b/tensorflow/python/debug/lib/debug_utils_test.py
@@ -185,7 +185,7 @@
self.assertEqual(["file:///tmp/tfdbg_1", "file:///tmp/tfdbg_2"],
watch_0.debug_urls)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_allNodes(self):
debug_utils.watch_graph(
self._run_options,
@@ -217,7 +217,7 @@
self.assertTrue("p1" in node_names)
self.assertTrue("s" in node_names)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_nodeNameWhitelist(self):
debug_utils.watch_graph(
self._run_options,
@@ -232,7 +232,7 @@
sorted(["a1_init", "a1", "a1/Assign", "a1/read", "p1"]),
sorted(node_names))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_opTypeWhitelist(self):
debug_utils.watch_graph(
self._run_options,
@@ -258,7 +258,7 @@
["DebugIdentity"], ["file:///tmp/tfdbg_1"])
self.assertEqual(["p1"], node_names)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_tensorDTypeWhitelist(self):
debug_utils.watch_graph(
self._run_options,
@@ -271,7 +271,7 @@
["DebugIdentity"], ["file:///tmp/tfdbg_1"])
self.assertItemsEqual(["a1", "a1/Assign", "b", "b/Assign"], node_names)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_nodeNameAndTensorDTypeWhitelists(self):
debug_utils.watch_graph(
self._run_options,
@@ -285,7 +285,7 @@
["DebugIdentity"], ["file:///tmp/tfdbg_1"])
self.assertItemsEqual(["a1", "a1/Assign"], node_names)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_nodeNameBlacklist(self):
debug_utils.watch_graph_with_blacklists(
self._run_options,
@@ -300,7 +300,7 @@
sorted(["b_init", "b", "b/Assign", "b/read", "c", "s"]),
sorted(node_names))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_opTypeBlacklist(self):
debug_utils.watch_graph_with_blacklists(
self._run_options,
@@ -313,7 +313,7 @@
["DebugIdentity"], ["file:///tmp/tfdbg_1"])
self.assertEqual(sorted(["p1", "s"]), sorted(node_names))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_nodeNameAndOpTypeBlacklists(self):
debug_utils.watch_graph_with_blacklists(
self._run_options,
@@ -327,7 +327,7 @@
["DebugIdentity"], ["file:///tmp/tfdbg_1"])
self.assertEqual(["s"], node_names)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_tensorDTypeBlacklists(self):
debug_utils.watch_graph_with_blacklists(
self._run_options,
@@ -344,7 +344,7 @@
self.assertNotIn("b/Assign", node_names)
self.assertIn("s", node_names)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWatchGraph_nodeNameAndTensorDTypeBlacklists(self):
debug_utils.watch_graph_with_blacklists(
self._run_options,
diff --git a/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py b/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py
index 74498c8..2405e29 100644
--- a/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py
+++ b/tensorflow/python/debug/lib/dist_session_debug_grpc_test.py
@@ -44,6 +44,7 @@
from tensorflow.python.platform import tf_logging
+@test_util.run_v1_only("b/120545219")
class DistributedSessionDebugTest(test_util.TensorFlowTestCase):
"""Test the debugging of distributed sessions."""
diff --git a/tensorflow/python/debug/lib/session_debug_file_test.py b/tensorflow/python/debug/lib/session_debug_file_test.py
index f5f9ba2..16ab815 100644
--- a/tensorflow/python/debug/lib/session_debug_file_test.py
+++ b/tensorflow/python/debug/lib/session_debug_file_test.py
@@ -34,6 +34,7 @@
from tensorflow.python.platform import googletest
+@test_util.run_v1_only("b/120545219")
class SessionDebugFileTest(session_debug_testlib.SessionDebugTestBase):
def _debug_urls(self, run_number=None):
@@ -45,7 +46,6 @@
else:
return os.path.join(self._dump_root, "run_%d" % run_number)
- @test_util.run_deprecated_v1
def testAllowsDifferentWatchesOnDifferentRuns(self):
"""Test watching different tensors on different runs of the same graph."""
diff --git a/tensorflow/python/debug/lib/session_debug_grpc_test.py b/tensorflow/python/debug/lib/session_debug_grpc_test.py
index bfc9a3a..472e244 100644
--- a/tensorflow/python/debug/lib/session_debug_grpc_test.py
+++ b/tensorflow/python/debug/lib/session_debug_grpc_test.py
@@ -91,6 +91,7 @@
server.stop_server().wait()
+@test_util.run_v1_only("b/120545219")
class SessionDebugGrpcTest(session_debug_testlib.SessionDebugTestBase):
@classmethod
@@ -353,6 +354,7 @@
return urls
+@test_util.run_v1_only("b/120545219")
class SessionDebugGrpcGatingTest(test_util.TensorFlowTestCase):
"""Test server gating of debug ops."""
@@ -730,6 +732,7 @@
self.assertEqual("DebugNumericSummary", debug_watch.debug_op)
+@test_util.run_v1_only("b/120545219")
class DelayedDebugServerTest(test_util.TensorFlowTestCase):
def testDebuggedSessionRunWorksWithDelayedDebugServerStartup(self):
diff --git a/tensorflow/python/debug/lib/session_debug_testlib.py b/tensorflow/python/debug/lib/session_debug_testlib.py
index 25ef91b..5165feb 100644
--- a/tensorflow/python/debug/lib/session_debug_testlib.py
+++ b/tensorflow/python/debug/lib/session_debug_testlib.py
@@ -84,6 +84,7 @@
return (math_ops.multiply(self._w, input_), state)
+@test_util.run_v1_only("b/120545219")
class SessionDebugTestBase(test_util.TensorFlowTestCase):
"""Base class for unit tests of tfdbg running with tf.Session."""
diff --git a/tensorflow/python/debug/lib/source_utils_test.py b/tensorflow/python/debug/lib/source_utils_test.py
index 9083297..4f4aea0 100644
--- a/tensorflow/python/debug/lib/source_utils_test.py
+++ b/tensorflow/python/debug/lib/source_utils_test.py
@@ -216,6 +216,7 @@
os.remove(unrelated_source_path)
+@test_util.run_v1_only("b/120545219")
class ListSourceAgainstDumpTest(test_util.TensorFlowTestCase):
def createAndRunGraphWithWhileLoop(self):
diff --git a/tensorflow/python/debug/lib/stepper_test.py b/tensorflow/python/debug/lib/stepper_test.py
index 3839c67..9e78e20 100644
--- a/tensorflow/python/debug/lib/stepper_test.py
+++ b/tensorflow/python/debug/lib/stepper_test.py
@@ -33,6 +33,7 @@
from tensorflow.python.training import gradient_descent
+@test_util.run_v1_only("b/120545219")
class StepperTest(test_util.TensorFlowTestCase):
def setUp(self):
@@ -443,6 +444,7 @@
self.assertAllClose(-4.0, result["fz"]["z"])
+@test_util.run_v1_only("b/120545219")
class StepperTestWithPlaceHolders(test_util.TensorFlowTestCase):
def setUp(self):
@@ -577,6 +579,7 @@
self.assertAllClose([[-1.0], [6.0]], stepper.finalize())
+@test_util.run_v1_only("b/120545219")
class StepperAssignAddTest(test_util.TensorFlowTestCase):
def setUp(self):
@@ -692,6 +695,7 @@
self.assertAllClose(12.0, stepper.cont(self.v))
+@test_util.run_v1_only("b/120545219")
class StepperBackwardRunTest(test_util.TensorFlowTestCase):
def setUp(self):
diff --git a/tensorflow/python/debug/wrappers/disk_usage_test.py b/tensorflow/python/debug/wrappers/disk_usage_test.py
index 0874525..88b1cd5 100644
--- a/tensorflow/python/debug/wrappers/disk_usage_test.py
+++ b/tensorflow/python/debug/wrappers/disk_usage_test.py
@@ -32,6 +32,7 @@
from tensorflow.python.training import monitored_session
+@test_util.run_v1_only("b/120545219")
class DumpingDebugWrapperDiskUsageLimitTest(test_util.TensorFlowTestCase):
@classmethod
diff --git a/tensorflow/python/debug/wrappers/dumping_wrapper_test.py b/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
index 11011a5..42e3b09 100644
--- a/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/dumping_wrapper_test.py
@@ -41,6 +41,7 @@
from tensorflow.python.training import monitored_session
+@test_util.run_v1_only("b/120545219")
class DumpingDebugWrapperSessionTest(test_util.TensorFlowTestCase):
def setUp(self):
diff --git a/tensorflow/python/debug/wrappers/framework_test.py b/tensorflow/python/debug/wrappers/framework_test.py
index 68584b4..a50fa7c 100644
--- a/tensorflow/python/debug/wrappers/framework_test.py
+++ b/tensorflow/python/debug/wrappers/framework_test.py
@@ -141,6 +141,7 @@
return framework.OnRunEndResponse()
+@test_util.run_v1_only("b/120545219")
class DebugWrapperSessionTest(test_util.TensorFlowTestCase):
def _no_rewrite_session_config(self):
diff --git a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
index 149a749..e38df86 100644
--- a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
@@ -127,6 +127,7 @@
return e.exit_token
+@test_util.run_v1_only("b/120545219")
class LocalCLIDebugWrapperSessionTest(test_util.TensorFlowTestCase):
def setUp(self):
diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 2d9a176..887c61c 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -104,6 +104,18 @@
)
py_library(
+ name = "distribute",
+ srcs = [
+ "__init__.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":distribute_lib",
+ ":mirrored_strategy",
+ ],
+)
+
+py_library(
name = "distribute_lib",
srcs = [
"distribute_lib.py",
diff --git a/tensorflow/python/distribute/__init__.py b/tensorflow/python/distribute/__init__.py
new file mode 100644
index 0000000..4ff912a
--- /dev/null
+++ b/tensorflow/python/distribute/__init__.py
@@ -0,0 +1,25 @@
+# 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.
+# ==============================================================================
+"""Distribution Strategy library."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# pylint: disable=unused-import
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribution_strategy_context
+from tensorflow.python.distribute import mirrored_strategy
+# pylint: enable=unused-import
diff --git a/tensorflow/python/distribute/cluster_resolver/cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/cluster_resolver.py
index 7774ac0..73188bd 100644
--- a/tensorflow/python/distribute/cluster_resolver/cluster_resolver.py
+++ b/tensorflow/python/distribute/cluster_resolver/cluster_resolver.py
@@ -22,6 +22,8 @@
import six
+from tensorflow.python.client import session
+from tensorflow.python.framework import ops
from tensorflow.python.training.server_lib import ClusterSpec
@@ -32,6 +34,14 @@
return master
+def get_accelerator_devices(master, config_proto):
+ # TODO(frankchn): Add support for eager mode as well as graph mode.
+ with ops.Graph().as_default():
+ with session.Session(master, config=config_proto) as s:
+ devices = s.list_devices()
+ return devices
+
+
@six.add_metaclass(abc.ABCMeta)
class ClusterResolver(object):
"""Abstract class for all implementations of ClusterResolvers.
@@ -91,8 +101,11 @@
"""
raise NotImplementedError()
- @abc.abstractmethod
- def num_accelerators_per_worker(self, session_config=None):
+ def num_accelerators(self,
+ task_type=None,
+ task_index=None,
+ accelerator_type='GPU',
+ config_proto=None):
"""Returns the number of accelerator cores per worker.
This returns the number of accelerator cores (such as GPUs and TPUs)
@@ -100,11 +113,24 @@
should return 0. This method will query the master for this information
if it is not otherwise known.
+ Optionally, we allow callers to specify the task_type, task_index, and
+ rpc_layer, if they want to target a specific TensorFlow process to query
+ the number of accelerators. This is to support heterogenous environments,
+ where the number of accelerators cores per host is different.
+
Args:
- session_config: (Optional) Configuration for starting a new session to
+ task_type: (Optional) The type of the TensorFlow task of the machine we
+ want to query.
+ task_index: (Optional) The index of the TensorFlow task of the machine we
+ want to query.
+ accelerator_type: (Optional) The type of accelerator we are trying to
+ query (defaults to 'GPU').
+ config_proto: (Optional) Configuration for starting a new session to
query how many accelerator cores it has.
"""
- raise NotImplementedError()
+ master = self.master(task_type, task_index)
+ devices = get_accelerator_devices(master, config_proto)
+ return sum(1 for d in devices if d.device_type == accelerator_type)
@abc.abstractproperty
def environment(self):
@@ -116,7 +142,7 @@
"""Simple implementation of ClusterResolver that accepts a ClusterSpec."""
def __init__(self, cluster_spec, master='', task_type=None, task_index=None,
- environment='', num_accelerators_per_worker=0,
+ environment='', num_accelerators=0,
rpc_layer=None):
"""Creates a SimpleClusterResolver from a ClusterSpec."""
super(SimpleClusterResolver, self).__init__()
@@ -124,7 +150,7 @@
self._task_type = task_type
self._task_index = task_index
self._environment = environment
- self._num_accelerators_per_worker = num_accelerators_per_worker
+ self._num_accelerators = num_accelerators
self._rpc_layer = rpc_layer
if not isinstance(cluster_spec, ClusterSpec):
@@ -180,17 +206,27 @@
def environment(self):
return self._environment
- def num_accelerators_per_worker(self, session_config=None):
+ def num_accelerators(self,
+ task_type=None,
+ task_index=None,
+ accelerator_type='GPU',
+ config_proto=None):
"""Returns the number of accelerator cores per worker.
+ The SimpleClusterResolver does not do automatic detection of accelerators,
+ so a TensorFlow session will never be created, and thus all arguments are
+ unused and we simply return whatever was passed in when this object was
+ initialized.
+
Args:
- session_config: Unused. The SimpleClusterResolver does not do automatic
- detection of accelerators, so a TensorFlow session will never be
- created, and thus a `session_config` is never necessary here, and will
- be ignored.
+ task_type: Unused.
+ task_index: Unused.
+ accelerator_type: Unused.
+ config_proto: Unused.
"""
- del session_config
- return self._num_accelerators_per_worker
+ # Unused
+ del task_type, task_index, accelerator_type, config_proto
+ return self._num_accelerators
@property
def rpc_layer(self):
@@ -361,9 +397,13 @@
def environment(self):
return self._cluster_resolvers[0].environment
- def num_accelerators_per_worker(self, session_config=None):
- return self._cluster_resolvers[0].num_accelerators_per_worker(
- session_config)
+ def num_accelerators(self,
+ task_type=None,
+ task_index=None,
+ accelerator_type='GPU',
+ config_proto=None):
+ return self._cluster_resolvers[0].num_accelerators(
+ task_type, task_index, accelerator_type, config_proto)
@property
def rpc_layer(self):
diff --git a/tensorflow/python/distribute/cluster_resolver/cluster_resolver_test.py b/tensorflow/python/distribute/cluster_resolver/cluster_resolver_test.py
index b5448fa..0ff6b6b 100644
--- a/tensorflow/python/distribute/cluster_resolver/cluster_resolver_test.py
+++ b/tensorflow/python/distribute/cluster_resolver/cluster_resolver_test.py
@@ -18,11 +18,64 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.client import session
+from tensorflow.python.distribute.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver
from tensorflow.python.distribute.cluster_resolver import UnionClusterResolver
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
+mock = test.mock
+
+
+class MockBaseClusterResolver(ClusterResolver):
+
+ def cluster_spec(self):
+ return None
+
+ def master(self, task_type=None, task_index=None, rpc_layer=None):
+ return ""
+
+ def environment(self):
+ return ""
+
+
+class BaseClusterResolverTest(test.TestCase):
+
+ @mock.patch.object(session.BaseSession, "list_devices")
+ def testNumAcceleratorsSuccess(self, mock_list_devices):
+ device_names = [
+ "/job:worker/task:0/device:GPU:0",
+ "/job:worker/task:0/device:GPU:1",
+ "/job:worker/task:0/device:GPU:2",
+ "/job:worker/task:0/device:GPU:3",
+ ]
+ device_list = [
+ session._DeviceAttributes(
+ name, "GPU", 1024, 0) for name in device_names
+ ]
+ mock_list_devices.return_value = device_list
+
+ resolver = MockBaseClusterResolver()
+ self.assertEqual(resolver.num_accelerators(), 4)
+
+ @mock.patch.object(session.BaseSession, "list_devices")
+ def testNumAcceleratorsFilterSuccess(self, mock_list_devices):
+ device_names = [
+ "/job:worker/task:0/device:TPU:0",
+ "/job:worker/task:0/device:TPU:1",
+ "/job:worker/task:0/device:TPU:2",
+ "/job:worker/task:0/device:TPU:3",
+ ]
+ device_list = [
+ session._DeviceAttributes(
+ name, "TPU", 1024, 0) for name in device_names
+ ]
+ mock_list_devices.return_value = device_list
+
+ resolver = MockBaseClusterResolver()
+ self.assertEqual(resolver.num_accelerators(), 0)
+
class UnionClusterResolverTest(test.TestCase):
# TODO(frankchn): Transform to parameterized test after it is included in the
@@ -65,13 +118,13 @@
simple_resolver = SimpleClusterResolver(base_cluster_spec, task_type="ps",
task_index=1, environment="cloud",
- num_accelerators_per_worker=8,
+ num_accelerators=8,
rpc_layer="grpc")
self.assertEqual(simple_resolver.task_type, "ps")
self.assertEqual(simple_resolver.task_index, 1)
self.assertEqual(simple_resolver.environment, "cloud")
- self.assertEqual(simple_resolver.num_accelerators_per_worker(), 8)
+ self.assertEqual(simple_resolver.num_accelerators(), 8)
self.assertEqual(simple_resolver.rpc_layer, "grpc")
def testOverrideSimpleClusterResolver(self):
@@ -82,7 +135,7 @@
simple_resolver = SimpleClusterResolver(base_cluster_spec, task_type="ps",
task_index=1, environment="cloud",
- num_accelerators_per_worker=8,
+ num_accelerators=8,
rpc_layer="grpc")
simple_resolver.task_type = "worker"
@@ -130,7 +183,7 @@
})
resolver1 = SimpleClusterResolver(cluster_spec_1, task_type="ps",
task_index=1, environment="cloud",
- num_accelerators_per_worker=8,
+ num_accelerators=8,
rpc_layer="grpc")
cluster_spec_2 = server_lib.ClusterSpec({
@@ -139,7 +192,7 @@
})
resolver2 = SimpleClusterResolver(cluster_spec_2, task_type="worker",
task_index=2, environment="local",
- num_accelerators_per_worker=16,
+ num_accelerators=16,
rpc_layer="http")
union_resolver = UnionClusterResolver(resolver1, resolver2)
@@ -147,7 +200,7 @@
self.assertEqual(union_resolver.task_type, "ps")
self.assertEqual(union_resolver.task_index, 1)
self.assertEqual(union_resolver.environment, "cloud")
- self.assertEqual(union_resolver.num_accelerators_per_worker(), 8)
+ self.assertEqual(union_resolver.num_accelerators(), 8)
self.assertEqual(union_resolver.rpc_layer, "grpc")
union_resolver.task_type = "worker"
diff --git a/tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver.py
index b167bc8..0651261 100644
--- a/tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver.py
+++ b/tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver.py
@@ -51,7 +51,6 @@
task_type='worker',
task_index=0,
rpc_layer='grpc',
- num_accelerators_per_worker=0,
credentials='default',
service=None):
"""Creates a new GceClusterResolver object.
@@ -73,8 +72,6 @@
can be distinguished from each other.
rpc_layer: The RPC layer TensorFlow should use to communicate across
instances.
- num_accelerators_per_worker: Number of accelerators (GPUs) present per
- instance.
credentials: GCE Credentials. If nothing is specified, this defaults to
GoogleCredentials.get_application_default().
service: The GCE API object returned by the googleapiclient.discovery
@@ -200,7 +197,3 @@
@rpc_layer.setter
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
-
- def num_accelerators_per_worker(self, session_config=None):
- del session_config # Unused, since this is set manually in __init__.
- return self._num_accelerators_per_worker
diff --git a/tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver.py
index 041c081..7ff6ec0 100644
--- a/tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver.py
+++ b/tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver.py
@@ -18,7 +18,6 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.python.client import device_lib
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
from tensorflow.python.training import server_lib
@@ -108,16 +107,14 @@
Returns:
The name or URL of the session master.
"""
+ task_type = task_type if task_type is not None else self.task_type
+ task_index = task_index if task_index is not None else self.task_index
+
if task_type is not None and task_index is not None:
return format_master_url(
self.cluster_spec().task_address(task_type, task_index),
rpc_layer or self.rpc_layer)
- if self.task_type is not None and self.task_index is not None:
- return format_master_url(
- self.cluster_spec().task_address(self.task_type, self.task_index),
- rpc_layer or self.rpc_layer)
-
return ''
def cluster_spec(self):
@@ -167,7 +164,3 @@
on internal systems.
"""
return ''
-
- def num_accelerators_per_worker(self, session_config=None):
- local_devices = device_lib.list_local_devices(session_config)
- return len([d for d in local_devices if d.device_type == 'GPU'])
diff --git a/tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver.py
index fd3c6d6..9dbe25b 100644
--- a/tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver.py
+++ b/tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver.py
@@ -23,6 +23,7 @@
import subprocess
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
from tensorflow.python.training.server_lib import ClusterSpec
@@ -206,10 +207,13 @@
"""
task_type = task_type if task_type is not None else self.task_type
task_index = task_index if task_index is not None else self.task_index
- rpc_layer = rpc_layer or self.rpc_layer
- master = self.cluster_spec().task_address(task_type, task_index)
- return '%s://%s' % (rpc_layer, master) if rpc_layer else master
+ if task_type is not None and task_index is not None:
+ return format_master_url(
+ self.cluster_spec().task_address(task_type, task_index),
+ rpc_layer or self.rpc_layer)
+
+ return ''
@property
def environment(self):
@@ -221,6 +225,11 @@
"""
return ''
- def num_accelerators_per_worker(self, session_config=None):
- del session_config # Unused, since this is set in __init__ manually.
+ def num_accelerators(self,
+ task_type=None,
+ task_index=None,
+ accelerator_type='GPU',
+ config_proto=None):
+ # Unused, since this is set in __init__ manually.
+ del task_type, task_index, accelerator_type, config_proto
return self._gpus_per_node
diff --git a/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver.py
index a3246e7..8d530cc 100644
--- a/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver.py
+++ b/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver.py
@@ -54,8 +54,7 @@
task_type=None,
task_index=None,
rpc_layer=None,
- environment=None,
- num_accelerators_per_worker=0):
+ environment=None):
"""Creates a new TFConfigClusterResolver.
Args:
@@ -66,15 +65,11 @@
rpc_layer: (String, optional) Overrides the rpc layer TensorFlow uses.
environment: (String, optional) Overrides the environment TensorFlow
operates in.
- num_accelerators_per_worker: (Integer, optional) Specifies the number of
- accelerators (e.g. GPUs, TPUs, others) that each node has.
"""
-
self._task_type = task_type
self._task_index = task_index
self._rpc_layer = rpc_layer
self._environment = environment
- self._num_accelerators_per_worker = num_accelerators_per_worker
@property
def task_type(self):
@@ -115,11 +110,6 @@
def rpc_layer(self, rpc_layer):
self._rpc_layer = rpc_layer
- def num_accelerators_per_worker(self, session_config=None):
- # TODO(frankchn): Connect to server (w/ session_config) in the future.
- del session_config # Unused, we do not connect to another server here.
- return self._num_accelerators_per_worker
-
def cluster_spec(self):
"""Returns a ClusterSpec based on the TF_CONFIG environment variable.
diff --git a/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver_test.py b/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver_test.py
index c20e51b..36b3bb9 100644
--- a/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver_test.py
+++ b/tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver_test.py
@@ -168,13 +168,11 @@
}
"""
- cluster_resolver = TFConfigClusterResolver(task_type='ps', task_index=0,
- num_accelerators_per_worker=8)
+ cluster_resolver = TFConfigClusterResolver(task_type='ps', task_index=0)
self.assertEqual('grpc://ps0:2222', cluster_resolver.master())
self.assertEqual('ps', cluster_resolver.task_type)
self.assertEqual(0, cluster_resolver.task_index)
- self.assertEqual(8, cluster_resolver.num_accelerators_per_worker())
cluster_resolver.task_type = 'worker'
cluster_resolver.task_index = 1
diff --git a/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py
index 1956bd7..52ac07d 100644
--- a/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py
+++ b/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py
@@ -18,13 +18,18 @@
from __future__ import division
from __future__ import print_function
+import collections
import os
+import re
from six.moves.urllib.request import Request
from six.moves.urllib.request import urlopen
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import get_accelerator_devices
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
from tensorflow.python.util import compat
@@ -41,6 +46,45 @@
_DEFAULT_ENV_VARIABLE = 'TPU_NAME'
_DISCOVERY_SERVICE_URL_ENV_VARIABLE = 'TPU_API_DISCOVERY_URL'
+_TPU_DEVICE_REGEX = re.compile(
+ r'.*task:(?P<host_id>\d+)/.*device:TPU:(?P<core_id>\d+)$')
+_TPU_CONN_RETRIES = 120
+
+DeviceDetails = collections.namedtuple(
+ 'DeviceDetails', ['device_map', 'total_cores'])
+
+
+def _get_device_dict_and_cores(devices):
+ """Returns a dict of hosts to cores and total cores given devices names.
+
+ Returns a namedtuple with two attributes:
+ device_map: A map of host_ids to a list of core_ids.
+ total_cores: The total number of cores within the TPU system.
+
+ Args:
+ devices: A list of devices returned by session.list_devices()
+ """
+ device_map = collections.defaultdict(list)
+ num_cores = 0
+ for device in devices:
+ match = _TPU_DEVICE_REGEX.match(device.name)
+ if match:
+ host_id = match.group('host_id')
+ core_id = match.group('core_id')
+ device_map[host_id].append(core_id)
+ num_cores += 1
+ return DeviceDetails(device_map, num_cores)
+
+
+def _verify_and_return_same_core_count(device_dict):
+ """Verifies that every device in device_dict has the same number of cores."""
+ num_cores_per_host_set = (
+ {len(core_ids) for core_ids in device_dict.values()})
+ if len(num_cores_per_host_set) != 1:
+ raise RuntimeError('TPU cores on each device is not the same. This '
+ 'should never happen. Devices: {}'.format(device_dict))
+ return num_cores_per_host_set.pop()
+
class TPUClusterResolver(ClusterResolver):
"""Cluster Resolver for Google Cloud TPUs.
@@ -92,7 +136,8 @@
self._tpu == compat.as_bytes('local') or
self._tpu.startswith(compat.as_bytes('/bns')) or
self._tpu.startswith(compat.as_bytes('localhost:')) or
- self._tpu.startswith(compat.as_bytes('grpc://'))):
+ self._tpu.startswith(compat.as_bytes('grpc://')) or
+ self._tpu.startswith(compat.as_bytes('uptc://'))):
return False
return True
@@ -196,13 +241,14 @@
elif tpu == 'local' or not tpu:
# Google environment, where the TPU is attached to the host.
self._environment = 'google'
- elif tpu.startswith('/bns'):
+ elif tpu.startswith('/bns') or tpu.startswith('uptc://'):
# Google environment, where we reach the TPU through BNS.
self._environment = 'google'
# If TPU is in the Google environment or exists locally, we don't use any
# RPC layer.
- if tpu.startswith('/bns') or tpu == 'local' or not tpu:
+ if tpu.startswith('/bns') or tpu.startswith(
+ 'uptc://') or tpu == 'local' or not tpu:
self.rpc_layer = None
else:
self.rpc_layer = 'grpc'
@@ -385,18 +431,49 @@
return server_lib.ClusterSpec(cluster_spec)
- def num_accelerators_per_worker(self, session_config=None):
+ def num_accelerators(self,
+ task_type=None,
+ task_index=None,
+ accelerator_type='TPU',
+ config_proto=None):
"""Returns the number of TPU cores per worker.
- This defaults to 8 for all current TPU configurations, and we do not need
- to query any remote systems for this.
+ Connects to the master and list all the devices present in the master,
+ and counts them up. Also verifies that the device counts per host in the
+ cluster is the same before returning the number of TPU cores per host.
Args:
- session_config: Unused. Not currently necessary to query anything as this
- number is 8 for all TPU configurations.
+ task_type: Unused.
+ task_index: Unused.
+ accelerator_type: Unused.
+ config_proto: Used to create a connection to a TPU master in order to
+ retrieve the system metadata.
+
+ Raises:
+ RuntimeError: If we cannot talk to a TPU worker after retrying or if the
+ number of TPU devices per host is different.
"""
- del session_config # Unused. Not necessary to query anything.
- return 8
+ retry_count = 1
+ # TODO(b/120564445): Replace with standard library for retries.
+ while True:
+ try:
+ device_details = _get_device_dict_and_cores(
+ get_accelerator_devices(self.master(), config_proto=config_proto))
+ break
+ except errors.DeadlineExceededError:
+ error_message = ('Failed to connect to master. The TPU might not be '
+ 'ready (e.g. still scheduling) or the master '
+ 'address is incorrect: got (%s)' % self.master())
+ if retry_count <= _TPU_CONN_RETRIES:
+ logging.warning(error_message)
+ logging.warning('Retrying (%d/%d)...', retry_count, _TPU_CONN_RETRIES)
+ retry_count += 1
+ else:
+ raise RuntimeError(error_message)
+
+ if device_details.total_cores:
+ return _verify_and_return_same_core_count(device_details.device_map)
+ return 0
@property
def environment(self):
@@ -404,7 +481,8 @@
return self._environment
def _start_local_server(self):
- address = self._requestComputeMetadata('instance/network-interfaces/0/ip')
+ address = compat.as_text(self._requestComputeMetadata(
+ 'instance/network-interfaces/0/ip'))
self._server = server_lib.Server(
{
'local': ['0.0.0.0:0']
diff --git a/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver_test.py b/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver_test.py
index 0f22ede..27d9260 100644
--- a/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver_test.py
+++ b/tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver_test.py
@@ -20,7 +20,10 @@
import os
-from tensorflow.python.distribute.cluster_resolver import TPUClusterResolver
+from tensorflow.python.client import session
+from tensorflow.python.distribute import cluster_resolver
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.framework import errors
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
from tensorflow.python.util import compat
@@ -101,7 +104,8 @@
return mock_client
- @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata',
+ @mock.patch.object(cluster_resolver.TPUClusterResolver,
+ '_requestComputeMetadata',
mock_request_compute_metadata)
def testRetrieveProjectAndZoneFromMetadata(self):
tpu_map = {
@@ -112,7 +116,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project=None,
zone=None,
tpu=['test-tpu-1'],
@@ -120,7 +124,7 @@
service=self.mock_service_client(tpu_map=tpu_map),
coordinator_name='coordinator')
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job {
name: 'coordinator'
@@ -130,11 +134,12 @@
name: 'worker'
tasks { key: 0 value: '10.1.2.3:8470' }
}
- """ % tpu_cluster_resolver._coordinator_port
+ """ % resolver._coordinator_port
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+ self.assertEqual(resolver.master(), 'grpc://10.1.2.3:8470')
- @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata',
+ @mock.patch.object(cluster_resolver.TPUClusterResolver,
+ '_requestComputeMetadata',
mock_request_compute_metadata)
def testRetrieveProjectAndZoneFromMetadataNoCoordinator(self):
tpu_map = {
@@ -145,7 +150,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project=None,
zone=None,
tpu=['test-tpu-1'],
@@ -153,14 +158,15 @@
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+ self.assertEqual(resolver.master(), 'grpc://10.1.2.3:8470')
- @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata',
+ @mock.patch.object(cluster_resolver.TPUClusterResolver,
+ '_requestComputeMetadata',
mock_request_compute_metadata)
def testUnhealthyCloudTpu(self):
tpu_map = {
@@ -171,7 +177,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project=None,
zone=None,
tpu='test-tpu-1',
@@ -180,9 +186,10 @@
service=self.mock_service_client(tpu_map=tpu_map))
with self.assertRaises(RuntimeError):
- tpu_cluster_resolver.cluster_spec()
+ resolver.cluster_spec()
- @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata',
+ @mock.patch.object(cluster_resolver.TPUClusterResolver,
+ '_requestComputeMetadata',
mock_request_compute_metadata)
def testNotReadyCloudTpu(self):
tpu_map = {
@@ -193,7 +200,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project=None,
zone=None,
tpu='test-tpu-1',
@@ -202,7 +209,7 @@
service=self.mock_service_client(tpu_map=tpu_map))
with self.assertRaises(RuntimeError):
- tpu_cluster_resolver.cluster_spec()
+ resolver.cluster_spec()
def testSimpleSuccessfulRetrieval(self):
tpu_map = {
@@ -213,7 +220,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu=['test-tpu-1'],
@@ -222,13 +229,13 @@
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } }
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+ self.assertEqual(resolver.master(), 'grpc://10.1.2.3:8470')
def testNewNetworkEndpointFormat(self):
tpu_map = {
@@ -241,7 +248,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
@@ -250,15 +257,16 @@
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } }
job { name: 'worker' tasks { key: 0 value: '10.2.3.4:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
- self.assertEqual('grpc://10.2.3.4:8470', tpu_cluster_resolver.master())
+ self.assertEqual('grpc://10.2.3.4:8470', resolver.master())
- @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata',
+ @mock.patch.object(cluster_resolver.TPUClusterResolver,
+ '_requestComputeMetadata',
mock_request_compute_metadata)
def testPodResolution(self):
tpu_map = {
@@ -286,13 +294,13 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
tpu='test-tpu-1',
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map),
coordinator_name='coordinator')
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job {
name: 'coordinator',
@@ -305,9 +313,9 @@
tasks { key: 2 value: '10.2.3.6:8470' }
tasks { key: 3 value: '10.2.3.7:8470' }
}
- """ % tpu_cluster_resolver._coordinator_port
+ """ % resolver._coordinator_port
self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.2.3.4:8470')
+ self.assertEqual(resolver.master(), 'grpc://10.2.3.4:8470')
def testPodResolutionNoCoordinator(self):
tpu_map = {
@@ -335,7 +343,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
@@ -343,7 +351,7 @@
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
@@ -354,13 +362,13 @@
}
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.2.3.4:8470')
+ self.assertEqual(resolver.master(), 'grpc://10.2.3.4:8470')
def testGetMasterNoEntries(self):
tpu_map = {}
with self.assertRaises(ValueError):
- TPUClusterResolver(
+ cluster_resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu=[],
@@ -370,14 +378,14 @@
# TODO(saeta): Convert to parameterized test when included in OSS TF.
def verifyShouldResolve(self, tpu, should_resolve):
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu=tpu,
coordinator_name=None,
credentials=None,
service=self.mock_service_client(tpu_map={}))
- self.assertEqual(should_resolve, tpu_cluster_resolver._shouldResolve(),
+ self.assertEqual(should_resolve, resolver._shouldResolve(),
"TPU: '%s'" % tpu)
def testShouldResolveNoName(self):
@@ -402,25 +410,26 @@
self.verifyShouldResolve('grpctpu', True)
def testNoCallComputeMetadata(self):
- tpu_cluster_resolver = TPUClusterResolver(tpu='/bns/foo/bar')
+ resolver = cluster_resolver.TPUClusterResolver(
+ tpu='/bns/foo/bar')
self.assertEqual(
- compat.as_bytes('/bns/foo/bar'), tpu_cluster_resolver.master())
- self.assertEqual(None, tpu_cluster_resolver.cluster_spec())
+ compat.as_bytes('/bns/foo/bar'), resolver.master())
+ self.assertEqual(None, resolver.cluster_spec())
def testGkeEnvironmentForDonut(self):
os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470'
self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ)
- self.assertTrue(TPUClusterResolver._inGke())
+ self.assertTrue(cluster_resolver.TPUClusterResolver._inGke())
self.assertEqual(
compat.as_bytes('grpc://10.120.27.5:8470'),
- compat.as_bytes(TPUClusterResolver._gkeEndpoints()))
+ compat.as_bytes(cluster_resolver.TPUClusterResolver._gkeEndpoints()))
- tpu_cluster_resolver = TPUClusterResolver()
+ resolver = cluster_resolver.TPUClusterResolver()
self.assertEqual(
compat.as_bytes('grpc://10.120.27.5:8470'),
- compat.as_bytes(tpu_cluster_resolver.master()))
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ compat.as_bytes(resolver.master()))
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
@@ -438,19 +447,19 @@
'grpc://10.120.27.8:8470')
self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ)
- self.assertTrue(TPUClusterResolver._inGke())
+ self.assertTrue(cluster_resolver.TPUClusterResolver._inGke())
self.assertEqual(
compat.as_bytes('grpc://10.120.27.5:8470,'
'grpc://10.120.27.6:8470,'
'grpc://10.120.27.7:8470,'
'grpc://10.120.27.8:8470'),
- compat.as_bytes(TPUClusterResolver._gkeEndpoints()))
+ compat.as_bytes(cluster_resolver.TPUClusterResolver._gkeEndpoints()))
- tpu_cluster_resolver = TPUClusterResolver()
+ resolver = cluster_resolver.TPUClusterResolver()
self.assertEqual(
compat.as_bytes('grpc://10.120.27.5:8470'),
- compat.as_bytes(tpu_cluster_resolver.master()))
- actual_cluster_spec = tpu_cluster_resolver.cluster_spec()
+ compat.as_bytes(resolver.master()))
+ actual_cluster_spec = resolver.cluster_spec()
expected_proto = """
job {
name: 'worker'
@@ -467,18 +476,21 @@
def testEnvironmentDiscoveryUrl(self):
os.environ['TPU_API_DISCOVERY_URL'] = 'https://{api}.internal/{apiVersion}'
self.assertEqual('https://{api}.internal/{apiVersion}',
- TPUClusterResolver._environmentDiscoveryUrl())
+ (cluster_resolver.TPUClusterResolver.
+ _environmentDiscoveryUrl()))
def testEnvironmentAndRpcDetectionForGoogle(self):
- tpu_cluster_resolver = TPUClusterResolver(tpu='/bns/ab/cd/ef')
- self.assertEqual(tpu_cluster_resolver.environment, 'google')
- self.assertEqual(tpu_cluster_resolver.rpc_layer, None)
+ resolver = cluster_resolver.TPUClusterResolver(
+ tpu='/bns/ab/cd/ef')
+ self.assertEqual(resolver.environment, 'google')
+ self.assertEqual(resolver.rpc_layer, None)
def testEnvironmentAndRpcDetectionForGrpcString(self):
- tpu_cluster_resolver = TPUClusterResolver(tpu='grpc://10.1.2.3:8470')
- self.assertEqual(tpu_cluster_resolver.environment, '')
- self.assertEqual(tpu_cluster_resolver.rpc_layer, 'grpc')
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+ resolver = cluster_resolver.TPUClusterResolver(
+ tpu='grpc://10.1.2.3:8470')
+ self.assertEqual(resolver.environment, '')
+ self.assertEqual(resolver.rpc_layer, 'grpc')
+ self.assertEqual(resolver.master(), 'grpc://10.1.2.3:8470')
def testOverrideTaskTypeAndIndexAndGetMaster(self):
tpu_map = {
@@ -506,7 +518,7 @@
}
}
- tpu_cluster_resolver = TPUClusterResolver(
+ resolver = cluster_resolver.TPUClusterResolver(
project='test-project',
zone='us-central1-c',
tpu='test-tpu-1',
@@ -514,17 +526,103 @@
credentials=None,
service=self.mock_service_client(tpu_map=tpu_map))
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.2.3.4:8470')
+ self.assertEqual(resolver.master(), 'grpc://10.2.3.4:8470')
- tpu_cluster_resolver.task_type = 'worker'
- tpu_cluster_resolver.task_index = 3
- self.assertEqual(tpu_cluster_resolver.master(), 'grpc://10.2.3.7:8470')
+ resolver.task_type = 'worker'
+ resolver.task_index = 3
+ self.assertEqual(resolver.master(), 'grpc://10.2.3.7:8470')
self.assertEqual(
- tpu_cluster_resolver.master(
+ resolver.master(
task_type='worker', task_index=2, rpc_layer='test'),
'test://10.2.3.6:8470')
+ def testGetDeviceDictAndCoresWithTPUs(self):
+ device_names = [
+ '/job:tpu_worker/task:0/device:TPU:0',
+ '/job:tpu_worker/task:1/device:TPU:1',
+ '/job:tpu_worker/task:2/device:TPU:0',
+ '/job:tpu_worker/task:3/device:TPU:1',
+ '/job:tpu_worker/task:0/device:TPU:4',
+ '/job:tpu_worker/task:1/device:TPU:5',
+ '/job:tpu_worker/task:2/device:TPU:4',
+ '/job:tpu_worker/task:3/device:TPU:5',
+ ]
+ device_list = [
+ session._DeviceAttributes(
+ name, 'TPU', 1024, 0) for name in device_names
+ ]
+
+ device_details = tpu_cluster_resolver._get_device_dict_and_cores(
+ device_list)
+ self.assertEqual(device_details.total_cores, 8)
+ self.assertEqual(device_details.device_map,
+ {'0': ['0', '4'],
+ '1': ['1', '5'],
+ '2': ['0', '4'],
+ '3': ['1', '5']})
+
+ def testGetDeviceDictAndCoresWithCPUsAndGPUs(self):
+ device_names = [
+ '/job:tpu_worker/task:0/device:CPU:0',
+ '/job:tpu_worker/task:1/device:CPU:0',
+ '/job:tpu_worker/task:2/device:CPU:0',
+ '/job:tpu_worker/task:3/device:CPU:0',
+ '/job:tpu_worker/task:0/device:GPU:1',
+ '/job:tpu_worker/task:1/device:GPU:1',
+ '/job:tpu_worker/task:2/device:GPU:1',
+ '/job:tpu_worker/task:3/device:GPU:1',
+ ]
+ device_list = [
+ session._DeviceAttributes(
+ name, 'XLA', 1024, 0) for name in device_names
+ ]
+
+ device_dict, num_cores = tpu_cluster_resolver._get_device_dict_and_cores(
+ device_list)
+ self.assertEqual(num_cores, 0)
+ self.assertEqual(device_dict, {})
+
+ def testVerifySameCoreCount(self):
+ self.assertEqual(
+ tpu_cluster_resolver._verify_and_return_same_core_count(
+ {0: [0, 1, 2, 3, 4, 5, 6, 7]}), 8)
+ self.assertEqual(
+ tpu_cluster_resolver._verify_and_return_same_core_count(
+ {0: [0, 1], 1: [2, 3]}), 2)
+ with self.assertRaises(RuntimeError):
+ tpu_cluster_resolver._verify_and_return_same_core_count(
+ {0: [0], 1: [1, 2]})
+
+ @mock.patch.object(session.BaseSession, 'list_devices')
+ def testNumAcceleratorsSuccess(self, mock_list_devices):
+ device_names = [
+ '/job:tpu_worker/task:0/device:TPU:0',
+ '/job:tpu_worker/task:1/device:TPU:1',
+ '/job:tpu_worker/task:2/device:TPU:0',
+ '/job:tpu_worker/task:3/device:TPU:1',
+ '/job:tpu_worker/task:0/device:TPU:4',
+ '/job:tpu_worker/task:1/device:TPU:5',
+ '/job:tpu_worker/task:2/device:TPU:4',
+ '/job:tpu_worker/task:3/device:TPU:5',
+ ]
+ device_list = [
+ session._DeviceAttributes(
+ name, 'TPU', 1024, 0) for name in device_names
+ ]
+ mock_list_devices.return_value = device_list
+
+ resolver = cluster_resolver.TPUClusterResolver(tpu='')
+ self.assertEqual(resolver.num_accelerators(), 2)
+
+ @mock.patch.object(session.BaseSession, 'list_devices')
+ def testNumAcceleratorsRetryFailure(self, mock_list_devices):
+ resolver = cluster_resolver.TPUClusterResolver(tpu='')
+ mock_list_devices.side_effect = errors.DeadlineExceededError(
+ None, None, 'timeout')
+ with self.assertRaises(RuntimeError):
+ resolver.num_accelerators()
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/distribute/cross_device_ops.py b/tensorflow/python/distribute/cross_device_ops.py
index a88ed62..6bb3639 100644
--- a/tensorflow/python/distribute/cross_device_ops.py
+++ b/tensorflow/python/distribute/cross_device_ops.py
@@ -53,15 +53,52 @@
if not isinstance(
destinations,
(value_lib.DistributedValues, resource_variable_ops.ResourceVariable,
- value_lib.AggregatingVariable, six.string_types, list)):
+ value_lib.AggregatingVariable, six.string_types, list, tuple)):
raise ValueError("destinations must be one of a `DistributedValues` object,"
- " a tf.Variable object, a device string, a list of device "
- "strings")
+ " a tf.Variable object, a device string, a list or tuple "
+ "of device strings")
if not check_destinations(destinations):
raise ValueError("destinations can not be empty")
+def reduce_non_distributed_value(extended, reduce_op, value, destinations):
+ """Reduce a non-DistributedValue `value` to `destinations`."""
+ if isinstance(value, value_lib.DistributedValues):
+ raise ValueError("You are passing a `DistributedValue` to "
+ "`reduce_non_distributed_value`, which is not allowed.")
+
+ # If the same value is present on all replicas then the PerReplica value will
+ # be a single value. We also handle the case when `value` is a single value
+ # and equal to 0.
+ if value == 0:
+ return 0
+ # If there is only a single value and the reduce op is MEAN,
+ # that value should be on all destinations.
+ if reduce_op == reduce_util.ReduceOp.MEAN:
+ return value
+
+ validate_destinations(destinations)
+ # We do not support a reduce op of SUM if the value is the same across
+ # all replicas. We call this as part of assign functions for MirroredVariables
+ # and summing up identical values across replicas is not clearly defined.
+ if (len(extended.worker_devices) != 1 or
+ not check_destinations(destinations)):
+ raise ValueError("A non-DistributedValues value %s cannot be reduced with "
+ "the given reduce op %s." % (value, reduce_op))
+ # TODO(anjalisridhar): Moves these methods to a device utility file?
+ devices = get_devices_from(destinations)
+ if len(devices) == 1:
+ with ops.device(devices[0]):
+ return array_ops.identity(value)
+ else:
+ value_updates = {}
+ for d in devices:
+ with ops.device(d):
+ value_updates[d] = array_ops.identity(value)
+ return value_lib.Mirrored(value_updates)
+
+
def _make_tensor_into_per_replica(input_tensor):
"""Converts a single tensor into a PerReplica object."""
if isinstance(input_tensor, (tuple, list)):
diff --git a/tensorflow/python/distribute/distribute_coordinator.py b/tensorflow/python/distribute/distribute_coordinator.py
index c0f9b8a..78c995a 100644
--- a/tensorflow/python/distribute/distribute_coordinator.py
+++ b/tensorflow/python/distribute/distribute_coordinator.py
@@ -29,6 +29,7 @@
from tensorflow.python.distribute import distribute_coordinator_context
from tensorflow.python.distribute import multi_worker_util
from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import coordinator
from tensorflow.python.training import monitored_session
from tensorflow.python.training import server_lib
@@ -328,7 +329,8 @@
task_id,
session_config,
rpc_layer="",
- worker_barrier=None):
+ worker_barrier=None,
+ coord=None):
"""Runs a single worker by calling `worker_fn` under context."""
session_config = copy.deepcopy(session_config)
strategy = copy.deepcopy(strategy)
@@ -350,7 +352,11 @@
rpc_layer=rpc_layer,
worker_barrier=worker_barrier)
with context:
- return worker_fn(strategy)
+ if coord:
+ with coord.stop_on_exception():
+ return worker_fn(strategy)
+ else:
+ return worker_fn(strategy)
def _split_cluster_for_evaluator(cluster_spec, task_type):
@@ -423,6 +429,7 @@
def _run_between_graph_client(worker_fn, strategy, eval_fn, eval_strategy,
cluster_spec, session_config, rpc_layer):
"""Runs a standalone client for between-graph replication."""
+ coord = coordinator.Coordinator()
eval_thread = None
if _TaskType.EVALUATOR in cluster_spec.jobs:
eval_thread = threading.Thread(
@@ -431,6 +438,7 @@
session_config),
kwargs={
"rpc_layer": rpc_layer,
+ "coord": coord,
})
eval_thread.start()
@@ -444,18 +452,18 @@
session_config),
kwargs={
"rpc_layer": rpc_layer,
- "worker_barrier": worker_barrier
+ "worker_barrier": worker_barrier,
+ "coord": coord,
})
t.start()
threads.append(t)
- # TODO(yuefengz): wrap threads into thread coordinator?
- for t in threads:
- t.join()
-
- # TODO(yuefengz): is it necessary to join eval thread?
if eval_thread:
- eval_thread.join()
+ # TODO(yuefengz): is it necessary to join eval thread?
+ threads_to_join = threads + [eval_thread]
+ else:
+ threads_to_join = threads
+ coord.join(threads_to_join)
# TODO(yuefengz): we probably want to return results from all workers?
return None
@@ -464,6 +472,7 @@
def _run_in_graph_client(worker_fn, strategy, eval_fn, eval_strategy,
cluster_spec, session_config, rpc_layer):
"""Runs a standalone client for in-graph replication."""
+ coord = coordinator.Coordinator()
eval_thread = None
if _TaskType.EVALUATOR in cluster_spec.jobs:
eval_thread = threading.Thread(
@@ -472,6 +481,7 @@
session_config),
kwargs={
"rpc_layer": rpc_layer,
+ "coord": coord,
})
eval_thread.start()
@@ -482,9 +492,12 @@
None,
None,
session_config,
- rpc_layer=rpc_layer)
+ rpc_layer=rpc_layer,
+ coord=coord)
+
if eval_thread:
- eval_thread.join()
+ coord.join([eval_thread])
+
return worker_result
diff --git a/tensorflow/python/distribute/distribute_coordinator_test.py b/tensorflow/python/distribute/distribute_coordinator_test.py
index f2cb950..dbed3e7 100644
--- a/tensorflow/python/distribute/distribute_coordinator_test.py
+++ b/tensorflow/python/distribute/distribute_coordinator_test.py
@@ -230,7 +230,7 @@
with ops.device("/job:worker/task:0"):
result = math_ops.add_n(xs)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
result_value = sess.run(result)
self.assertEqual(result_value, expected)
if result_value == expected:
@@ -278,7 +278,7 @@
train_op = control_flow_ops.group([x_add, y_sub])
if context.is_chief:
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Synchronize workers after initializaton.
if context.has_barrier:
@@ -427,6 +427,7 @@
# Each finished worker will increment self._result_correct.
self.assertEqual(self._result_correct, NUM_WORKERS)
+ @test_util.run_v1_only("b/120545219")
def testBetweenGraphWithMonitoredSession(self):
"""Test monitored session in standalone client mode."""
distribute_coordinator.run_distribute_coordinator(
@@ -600,6 +601,7 @@
# Each finished worker will increment self._result_correct.
self.assertEqual(self._result_correct, NUM_WORKERS)
+ @test_util.run_v1_only("b/120545219")
def testBetweenGraphWithMonitoredSession(self):
cluster_spec = self._create_cluster_spec(
num_workers=NUM_WORKERS, num_ps=NUM_PS)
diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index eddd6ff..60bb75d 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -422,6 +422,42 @@
return self.extended._make_input_fn_iterator( # pylint: disable=protected-access
input_fn, replication_mode=replication_mode)
+ def experimental_run(self, fn, input_iterator=None):
+ """Runs ops in `fn` on each replica, with inputs from `input_iterator`.
+
+ When eager execution is enabled, executes ops specified by `fn` on each
+ replica. Otherwise, builds a graph to execute the ops on each replica.
+
+ Each replica will take a single, different input from the inputs provided by
+ one `get_next` call on the input iterator.
+
+ `fn` may call `tf.distribute.get_replica_context()` to access members such
+ as `replica_id_in_sync_group`.
+
+ IMPORTANT: Depending on the `DistributionStrategy` being used, and whether
+ eager execution is enabled, `fn` may be called one or more times (once for
+ each replica).
+
+ Args:
+ fn: function to run. The inputs to the function must match the outputs of
+ `input_iterator.get_next()`. The output must be a `tf.nest` of
+ `Tensor`s.
+ input_iterator: (Optional) input iterator from which the inputs are taken.
+
+ Returns:
+ Merged return value of `fn` across replicas. The structure of the return
+ value is the same as the return value from `fn`. Each element in the
+ structure can either be `PerReplica` (if the values are unsynchronized),
+ `Mirrored` (if the values are kept in sync), or `Tensor` (if running on a
+ single replica).
+ """
+ with self.scope():
+ if input_iterator is None:
+ return self._extended.call_for_each_replica(fn)
+ else:
+ inputs = input_iterator.get_next()
+ return self._extended.call_for_each_replica(fn, args=(inputs,))
+
@doc_controls.do_not_generate_docs # DEPRECATED, moving to `extended`
def broadcast(self, tensor, destinations=None):
"""DEPRECATED: use extended.broadcast_to() instead."""
@@ -565,8 +601,8 @@
variable created in `scope`.
Returns:
- A list of values contained in `value`. If `value` represents a single
- value, this returns `[value].`
+ A tuple of values contained in `value`. If `value` represents a single
+ value, this returns `(value,).`
"""
return self._extended._unwrap(value) # pylint: disable=protected-access
@@ -1346,14 +1382,14 @@
@property
def worker_devices(self):
- """Returns the list of devices used to run `call_for_each_replica()` calls.
+ """Returns the tuple of all devices used to for compute replica execution.
"""
# TODO(josh11b): More docstring
raise NotImplementedError("must be implemented in descendants")
@property
def parameter_devices(self):
- """Returns the list of devices used for variable and `update` placement."""
+ """Returns the tuple of all devices used to place variables."""
# TODO(josh11b): More docstring
raise NotImplementedError("must be implemented in descendants")
@@ -1513,9 +1549,9 @@
@property
def devices(self):
- """The devices this replica is to be executed on, as a list of strings."""
+ """The devices this replica is to be executed on, as a tuple of strings."""
require_replica_context(self)
- return [device_util.current()]
+ return (device_util.current(),)
# TODO(josh11b): Implement `start_all_reduce(method, t)` for efficient
# all-reduce. It would return a function returning the result of reducing `t`
@@ -1605,7 +1641,7 @@
return array_ops.identity(replica_local_var)
def _unwrap(self, distributed_value):
- return [distributed_value]
+ return (distributed_value,)
def value_container(self, value):
return value
diff --git a/tensorflow/python/distribute/estimator_training.py b/tensorflow/python/distribute/estimator_training.py
index 549fa8f..7d5f231 100644
--- a/tensorflow/python/distribute/estimator_training.py
+++ b/tensorflow/python/distribute/estimator_training.py
@@ -354,7 +354,7 @@
if (estimator._config._distribute_coordinator_mode !=
dc.CoordinatorMode.STANDALONE_CLIENT):
raise ValueError('Only `STANDALONE_CLIENT` mode is supported when you call '
- '`Estimator.train`')
+ '`Estimator.evaluate`')
if estimator._config._eval_distribute.extended.experimental_between_graph:
# TODO(yuefengz): remove this limitation once we figure out how to merge
diff --git a/tensorflow/python/distribute/mirrored_strategy.py b/tensorflow/python/distribute/mirrored_strategy.py
index 0775920..fb3cf84 100644
--- a/tensorflow/python/distribute/mirrored_strategy.py
+++ b/tensorflow/python/distribute/mirrored_strategy.py
@@ -43,18 +43,24 @@
from tensorflow.python.ops import variable_scope
from tensorflow.python.training import coordinator
from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
# TODO(josh11b): Replace asserts in this file with if ...: raise ...
@contextlib.contextmanager
-def _enter_graph(g):
- if context.executing_eagerly():
+def _enter_graph(g, eager, creator_stack=None):
+ """Context manager for selecting a graph and maybe eager mode."""
+ if eager:
with g.as_default(), context.eager_mode():
+ if creator_stack is not None:
+ g._variable_creator_stack = creator_stack # pylint: disable=protected-access
yield
else:
with g.as_default():
+ if creator_stack is not None:
+ g._variable_creator_stack = creator_stack # pylint: disable=protected-access
yield
@@ -68,10 +74,9 @@
pass
-# _call_for_each_replica and _reduce_non_distributed_value are not members of
-# MirroredStrategy so that they are generally not allowed to use anything
-# specific to MirroredStrategy and thus can be shared with other distribution
-# strategies.
+# _call_for_each_replica is not a member of MirroredStrategy so that it is
+# not allowed to use anything specific to MirroredStrategy and thus
+# can be shared with other distribution strategies.
# TODO(yuefengz): maybe create a common class for those who need to call this
@@ -168,7 +173,12 @@
# capture the name_scope from the first MRT and assume it is
# the same for all other MRTs.
mtt_captured_name_scope = threads[0].captured_name_scope
- with ops.name_scope(mtt_captured_name_scope):
+ # Capture and merge the control dependencies from all the threads.
+ mtt_captured_control_deps = set()
+ for t in threads:
+ mtt_captured_control_deps.update(t.captured_control_deps)
+ with ops.name_scope(mtt_captured_name_scope),\
+ ops.control_dependencies(mtt_captured_control_deps):
merge_result = threads[0].merge_fn(distribution, *merge_args,
**merge_kwargs)
for t in threads:
@@ -181,43 +191,6 @@
return values.regroup({t.device: t.main_result for t in threads})
-def _reduce_non_distributed_value(extended, reduce_op, value, destinations):
- """Reduce a non-DistributedValue `value` to `destinations`."""
- if isinstance(value, values.DistributedValues):
- raise ValueError("You are passing a `DistributedValue` to "
- "`_reduce_non_distributed_value`, which is not allowed.")
-
- # If the same value is present on all replicas then the PerReplica value will
- # be a single value. We also handle the case when `value` is a single value
- # and equal to 0.
- if value == 0:
- return 0
- # If there is only a single value and the reduce op is MEAN,
- # that value should be on all destinations.
- if reduce_op == reduce_util.ReduceOp.MEAN:
- return value
-
- cross_device_ops_lib.validate_destinations(destinations)
- # We do not support a reduce op of SUM if the value is the same across
- # all replicas. We call this as part of assign functions for MirroredVariables
- # and summing up identical values across replicas is not clearly defined.
- if (len(extended.worker_devices) != 1 or
- not cross_device_ops_lib.check_destinations(destinations)):
- raise ValueError("A non-DistributedValues value %s cannot be reduced with "
- "the given reduce op %s." % (value, reduce_op))
- # TODO(anjalisridhar): Moves these methods to a device utility file?
- devices = cross_device_ops_lib.get_devices_from(destinations)
- if len(devices) == 1:
- with ops.device(devices[0]):
- return array_ops.identity(value)
- else:
- value_updates = {}
- for d in devices:
- with ops.device(d):
- value_updates[d] = array_ops.identity(value)
- return values.Mirrored(value_updates)
-
-
def _create_mirrored_variable(devices, real_mirrored_creator, *args, **kwargs): # pylint: disable=g-missing-docstring
# Figure out what collections this variable should be added to.
# We'll add the MirroredVariable to those collections instead.
@@ -422,6 +395,7 @@
("/device:CPU:0",))
+@tf_export("distribute.MirroredStrategy")
class MirroredStrategy(distribute_lib.DistributionStrategy):
"""Mirrors vars to distribute across multiple devices and machines.
@@ -470,7 +444,7 @@
assert len(set(devices)) == len(devices), (
"No duplicates allowed in `devices` argument.")
# TODO(josh11b): Require at least 2 devices?
- self._devices = [device_util.resolve(d) for d in devices]
+ self._devices = tuple(device_util.resolve(d) for d in devices)
self._canonical_device_set = set(self._devices)
self._device_index = values.PerReplica(
{d: i for i, d in enumerate(devices)})
@@ -486,7 +460,7 @@
assert len(set(devices)) == len(devices), (
"No duplicates allowed in `devices` argument.")
# TODO(josh11b): Require at least 2 devices?
- self._devices = [device_util.resolve(d) for d in devices]
+ self._devices = tuple(device_util.resolve(d) for d in devices)
self._canonical_device_set = set(self._devices)
self._device_index = values.PerReplica(
{d: i for i, d in enumerate(devices)})
@@ -693,14 +667,17 @@
return self._cross_device_ops or self._inferred_cross_device_ops
def _reduce_to(self, reduce_op, value, destinations):
+ if (isinstance(value, values.Mirrored) and
+ reduce_op == reduce_util.ReduceOp.MEAN):
+ return value
assert not isinstance(value, values.Mirrored)
if not isinstance(value, values.DistributedValues):
# This function handles reducing values that are not PerReplica or
# Mirrored values. For example, the same value could be present on all
# replicas in which case `value` would be a single value or value could
# be 0.
- return _reduce_non_distributed_value(self, reduce_op, value,
- destinations)
+ return cross_device_ops_lib.reduce_non_distributed_value(
+ self, reduce_op, value, destinations)
return self._get_cross_device_ops().reduce(
reduce_op, value, destinations=destinations)
@@ -722,7 +699,7 @@
return values.update_regroup(self, updates, group)
def _update_non_slot(self, colocate_with, fn, args, kwargs, group):
- assert isinstance(colocate_with, list)
+ assert isinstance(colocate_with, tuple)
# TODO(josh11b): In eager mode, use one thread per device.
updates = {}
for d in colocate_with:
@@ -743,9 +720,9 @@
if isinstance(val, values.DistributedValues):
# Return in a deterministic order.
if set(val.devices) == self._canonical_device_set:
- return [val.get(device=d) for d in self._devices]
- return [val.get(device=d) for d in sorted(val.devices)]
- return [val]
+ return tuple(val.get(device=d) for d in self._devices)
+ return tuple(val.get(device=d) for d in sorted(val.devices))
+ return (val,)
def value_container(self, val):
return values.value_container(val)
@@ -756,12 +733,11 @@
@property
def worker_devices(self):
- # Make a copy to prevent users from accidentally mutating our copy.
- return list(self._devices)
+ return self._devices
@property
def parameter_devices(self):
- return list(self._devices)
+ return self._devices
@property
def experimental_between_graph(self):
@@ -781,7 +757,7 @@
def non_slot_devices(self, var_list):
del var_list
- return list(self._devices)
+ return tuple(self._devices)
def _get_devices_from(self, colocate_with=None):
if colocate_with is None:
@@ -830,14 +806,19 @@
self.has_paused = threading.Event()
# These fields have to do with inheriting various contexts from the
# parent thread:
+ ctx = context.context()
+ self.in_eager = ctx.executing_eagerly()
# pylint: disable=protected-access
- self.context_mode = context.context()._eager_context.mode
- if not context.context()._context_handle:
- context.context()._initialize_handle_and_devices()
+ if not ctx._context_handle:
+ ctx._initialize_handle_and_devices()
self.context_device_policy = (
pywrap_tensorflow.TFE_ContextGetDevicePlacementPolicy(
- context.context()._context_handle))
+ ctx._context_handle))
self.graph = ops.get_default_graph()
+ with ops.init_scope():
+ self._init_in_eager = context.executing_eagerly()
+ self._init_graph = ops.get_default_graph()
+
self._variable_creator_stack = self.graph._variable_creator_stack[:]
self._captured_var_scope = variable_scope.get_variable_scope()
# Adding a "/" at end lets us re-enter this scope later.
@@ -851,16 +832,16 @@
def run(self):
# pylint: disable=protected-access
- self.graph._variable_creator_stack = self._variable_creator_stack
self.should_run.wait()
self.should_run.clear()
try:
if self.coord.should_stop():
return
with self.coord.stop_on_exception(), \
- context.context()._mode(self.context_mode), \
+ _enter_graph(self._init_graph, self._init_in_eager), \
+ _enter_graph(self.graph, self.in_eager,
+ self._variable_creator_stack), \
context.context().device_policy(self.context_device_policy), \
- _enter_graph(self.graph), \
MirroredReplicaContext(self.distribution, constant_op.constant(
self.replica_id, dtypes.int32)), \
ops.device(self.device), \
@@ -894,6 +875,8 @@
# Adding a "/" at end lets us re-enter this scope later.
if t.captured_name_scope:
t.captured_name_scope += "/"
+
+ t.captured_control_deps = t.graph._current_control_dependencies() # pylint: disable=protected-access
t.has_paused.set()
t.should_run.wait()
t.should_run.clear()
diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
index 01a1680..a5918b7 100644
--- a/tensorflow/python/distribute/values.py
+++ b/tensorflow/python/distribute/values.py
@@ -570,6 +570,12 @@
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
+ def __str__(self):
+ return "%s:%s" % (self.__class__.__name__, self._index)
+
+ def __repr__(self):
+ return "%s(%r)" % (self.__class__.__name__, self._index)
+
@property
def handle(self):
# If we're in a tpu.rewrite(), return the replicated handle.
diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index 29f9b2c..481f680 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -1104,8 +1104,13 @@
dimension of `target` and `source` do not match.
"""
target_shape = target.shape
- if not target_shape.with_rank_at_least(2)[0].is_compatible_with(
- source.shape.with_rank_at_least(2)[0]):
+ if target_shape.rank is None:
+ dim = Dimension(None)
+ else:
+ dim = target_shape.dims[0]
+ if not (target_shape.with_rank_at_least(2) and
+ source.shape.with_rank_at_least(2) and
+ dim.is_compatible_with(source.shape[0])):
raise ValueError(
"Need first dimension of target shape (%s) and "
"source shape (%s) to match." % (target.shape, source.shape))
diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py
index 3cec40a..477d18e 100644
--- a/tensorflow/python/eager/backprop_test.py
+++ b/tensorflow/python/eager/backprop_test.py
@@ -648,6 +648,7 @@
g.gradient(x, y)
@test_util.run_in_graph_and_eager_modes
+ @test_util.run_v1_only('b/120545219')
def testGradientTapeWithCond(self):
x = constant_op.constant(3.0)
@@ -669,6 +670,7 @@
self.assertEqual(self.evaluate(dy), 6.0)
@test_util.run_in_graph_and_eager_modes
+ @test_util.run_v1_only('b/120545219')
def testGradientTapeWithWhileLoop(self):
i = constant_op.constant(1)
x = constant_op.constant(2.)
@@ -704,6 +706,7 @@
@test_util.assert_no_new_tensors
@test_util.run_in_graph_and_eager_modes
+ @test_util.run_v1_only('b/120545219')
def testPersistentTape(self):
with backprop.GradientTape(persistent=True) as g:
x = constant_op.constant(3.0)
@@ -1243,16 +1246,19 @@
answer = [array_ops.diag(2 * x * y), array_ops.diag(x * x)]
return jacobian, answer
+ @test_util.run_v1_only('b/120545219')
def testPfor(self):
jacobian, answer = self._jacobian(experimental_use_pfor=True)
for j, a in zip(jacobian, answer):
self.assertAllEqual(a, j)
+ @test_util.run_v1_only('b/120545219')
def testWhileLoop(self):
jacobian, answer = self._jacobian(experimental_use_pfor=False)
for j, a in zip(jacobian, answer):
self.assertAllEqual(a, j)
+ @test_util.run_v1_only('b/120545219')
def testPforDefun(self):
@function.defun
@@ -1263,6 +1269,7 @@
for j, a in zip(jacobian, answer):
self.assertAllEqual(a, j)
+ @test_util.run_v1_only('b/120545219')
def testWhileLoopDefun(self):
@function.defun
@@ -1273,6 +1280,7 @@
for j, a in zip(jacobian, answer):
self.assertAllEqual(a, j)
+ @test_util.run_v1_only('b/120545219')
def testPersistentTape(self):
if not context.executing_eagerly():
return
@@ -1283,6 +1291,7 @@
with self.assertRaisesRegexp(RuntimeError, 'persistent'):
g.jacobian(y, x, experimental_use_pfor=False)
+ @test_util.run_v1_only('b/120545219')
def testPforException(self):
var = variables.Variable([1.])
@@ -1303,6 +1312,7 @@
with self.assertRaisesRegexp(ValueError, 'No converter'):
g.jacobian(y, x, experimental_use_pfor=True)
+ @test_util.run_v1_only('b/120545219')
def test_parallel_iterations(self):
with backprop.GradientTape(persistent=True) as g:
x = constant_op.constant([[1., 2], [3, 4]])
diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py
index 2f6b038..848b300 100644
--- a/tensorflow/python/eager/context.py
+++ b/tensorflow/python/eager/context.py
@@ -478,10 +478,6 @@
Raises:
ValueError: If name is not a string or is an invalid device name.
"""
- devices = self._context_devices
- if devices is None:
- self._initialize_handle_and_devices()
- devices = self._context_devices
eager_context = self._eager_context
old_device_name = eager_context.device_name
old_device_spec = eager_context.device_spec
@@ -502,7 +498,9 @@
if old_device_name:
new_device_spec = copy.copy(old_device_spec)
else:
- new_device_spec = pydev.DeviceSpec.from_string(devices[0])
+ self._initialize_handle_and_devices()
+ new_device_spec = pydev.DeviceSpec.from_string(
+ self._context_devices[0])
new_device_spec.merge_from(device_spec)
else:
new_device_spec = pydev.DeviceSpec.from_string("")
@@ -925,6 +923,10 @@
# but they do all import this file. Note that IS_IN_GRAPH_MODE and
# in_graph_mode are both parameterless functions.
def _tmp_in_graph_mode():
+ if context_safe() is None:
+ # Context not yet initialized. Assume graph mode following the
+ # default implementation in `is_in_graph_mode`.
+ return True
return not executing_eagerly()
diff --git a/tensorflow/python/eager/def_function.py b/tensorflow/python/eager/def_function.py
index 6bacd7a..a12f9ed 100644
--- a/tensorflow/python/eager/def_function.py
+++ b/tensorflow/python/eager/def_function.py
@@ -242,6 +242,7 @@
raise NotImplementedError()
self._created_variables = None
self._stateful_fn = None
+ self._stateless_fn = None
self._descriptor_cache = weakref.WeakKeyDictionary()
self._name = name
@@ -342,6 +343,10 @@
"""The python function wrapped in this tf.function."""
return self._python_function
+ @property
+ def input_signature(self):
+ return self._input_signature
+
def get_initialization_function(self, *args, **kwargs):
"""Returns a `Function` object which initializes this function's variables.
@@ -378,6 +383,26 @@
return initialize_variables.get_concrete_function()
+ @property
+ def _cached_input_signatures(self):
+ """All input signatures used to call this PolymorphicFunction."""
+ seen = set()
+ # Preserves signature ordering rather than returning a set() so that we
+ # don't need to re-sort signatures later to work around Python 2's set
+ # nondeterminism.
+ # pylint: disable=protected-access
+ concrete_functions = []
+ if self._stateful_fn:
+ concrete_functions.extend(self._stateful_fn._function_cache.values())
+ if self._stateless_fn:
+ concrete_functions.extend(self._stateless_fn._function_cache.values())
+ for concrete_function in concrete_functions:
+ signature = concrete_function._python_call_signature
+ if signature not in seen:
+ yield signature
+ seen.add(signature)
+ # pylint: enable=protected-access
+
def get_concrete_function(self, *args, **kwargs):
"""Returns a `Function` object specialized to inputs and execution context.
@@ -528,12 +553,33 @@
assert f(x, y).numpy() == g(x, y).numpy()
# Tensors and tf.Variables used by the Python function are captured in the
- # traced graph.
+ # graph.
@tf.function
def h():
return f(x, y)
assert (h().numpy() == f(x, y).numpy()).all()
+
+ # Data-dependent control flow is also captured in the graph. Supported
+ # control flow statements include `if`, `for`, `break`, `continue`, `return`.
+ @tf.function
+ def g(x):
+ if tf.reduce_sum(x) > 0:
+ return x * x
+ else:
+ return -x // 2
+
+ # print and TensorFlow side effects are supported, but exercise caution when
+ # using Python side effects like mutating objects, saving to files, etc.
+ l = []
+
+ @tf.function
+ def g(x):
+ for i in x:
+ print(i) # Works
+ tf.assign(v, i) # Works
+ tf.py_func(lambda i: l.append(i))(i) # Works
+ l.append(i) # Caution! Doesn't work.
```
_Referencing `tf.Variable`s_
@@ -605,6 +651,7 @@
```
_Input Signatures_
+
`function` instantiates a separate graph for every unique set of input
shapes and datatypes. For example, the following code snippet will result
in three distinct graphs being traced, as each input has a different
@@ -638,9 +685,15 @@
When an `input_signature` is specified, the callable will only accept `Tensor`
(or NumPy `ndarray`) objects as arguments.
- _Tracing_
- Note that `function` only traces TensorFlow operations, all the other
- Python code that `func` executes will shape the _construction_ of the graph.
+ _Tracing and staging_
+
+ When `autograph` is `True`, all Python code that depends on `Tensor` values is
+ staged into a TensorFlow graph. When `autograph` is `False`, the function is
+ traced and control flow is not allowed to depend on data.
+
+ Note that `function` only stages TensorFlow operations, all Python code that
+ `func` executes and does not depend on data will shape the _construction_ of
+ the graph.
For example, consider the following:
```python
@@ -653,21 +706,26 @@
```
`add_noise()` will return a different output every time it is invoked.
- However, `traced` will return the same value every time it is called, since a
- particular random value generated by the `np.random.randn` call will be
- inserted in the traced TensorFlow graph as a constant. In this particular
- example, replacing `np.random.randn(5, 5)` with `tf.random_normal((5, 5))`
- will result in the same behavior for `add_noise()` and `traced()`.
+ However, `add_noise` will return the same value every time it is called,
+ since a particular random value generated by the `np.random.randn` call will
+ be inserted in the traced/staged TensorFlow graph as a constant. In this
+ particular example, replacing `np.random.randn(5, 5)` with
+ `tf.random_normal((5, 5))` will result in the same behavior for `add_noise()`
+ and `traced()`.
_Python Side-Effects_
+
A corollary of the previous discussion on tracing is the following: If a
Python function `func` has Python side-effects, then executing `func` multiple
- times
- may not be semantically equivalent to executing `F = tf.function(func)`
+ times may not be semantically equivalent to executing `F = tf.function(func)`
multiple times; this difference is due to the fact that `function` only
captures the subgraph of TensorFlow operations that is constructed when `func`
is invoked to trace a graph.
+ The same is true if code with Python side effects is used inside control flow,
+ such as a loop. If your code uses side effects that are not intended to
+ control graph construction, wrap them inside `tf.py_func`.
+
Args:
func: function to be compiled. If `func` is None, returns a decorator that
can be invoked with a single argument - `func`. The end result is
diff --git a/tensorflow/python/eager/def_function_test.py b/tensorflow/python/eager/def_function_test.py
index 4100a10..8b4c407 100644
--- a/tensorflow/python/eager/def_function_test.py
+++ b/tensorflow/python/eager/def_function_test.py
@@ -238,6 +238,25 @@
concrete = compute.get_concrete_function(
tensor_spec.TensorSpec(None, dtypes.float32))
self.assertAllClose(4., concrete(constant_op.constant(2.)))
+ input_signature, = compute._cached_input_signatures
+ self.assertEqual(
+ tuple(input_signature),
+ (tensor_spec.TensorSpec(None, dtypes.float32),))
+
+ def test_serialization_signature_cache(self):
+
+ @def_function.function
+ def f(x, y):
+ return x, y
+
+ f(constant_op.constant([[3., 4.]]), constant_op.constant([2.]))
+ f(constant_op.constant([[3, 4, 5]]), constant_op.constant([2]))
+ self.assertEqual(
+ set(f._cached_input_signatures),
+ set(((tensor_spec.TensorSpec([1, 2], dtypes.float32),
+ tensor_spec.TensorSpec([1], dtypes.float32)),
+ (tensor_spec.TensorSpec([1, 3], dtypes.int32),
+ tensor_spec.TensorSpec([1], dtypes.int32)))))
if __name__ == '__main__':
diff --git a/tensorflow/python/eager/execute.py b/tensorflow/python/eager/execute.py
index 6f8c780..7415a0a 100644
--- a/tensorflow/python/eager/execute.py
+++ b/tensorflow/python/eager/execute.py
@@ -66,12 +66,6 @@
six.raise_from(core._status_to_exception(e.code, message), None)
except TypeError as e:
if any(ops._is_keras_symbolic_tensor(x) for x in inputs):
- if any(isinstance(x, ops.EagerTensor) for x in inputs):
- raise TypeError("You are attempting to mix computation of symbolic "
- "Tensors (computation rooted at tf.keras.Input()) "
- "and concrete values. This is not supported. "
- "If you need this support, file an issue on the "
- "TensorFlow GitHub repository.")
raise core._SymbolicException
raise e
# pylint: enable=protected-access
diff --git a/tensorflow/python/eager/execution_callbacks.py b/tensorflow/python/eager/execution_callbacks.py
index 28b6b84..af1afa3 100644
--- a/tensorflow/python/eager/execution_callbacks.py
+++ b/tensorflow/python/eager/execution_callbacks.py
@@ -20,6 +20,7 @@
import contextlib
import functools
+import enum # pylint: disable=g-bad-import-order
import numpy as np
@@ -29,13 +30,25 @@
from tensorflow.python.eager import execute
from tensorflow.python.platform import tf_logging as logging
-IGNORE = "ignore"
-PRINT = "print"
-RAISE = "raise"
-WARN = "warn"
-_DEFAULT_CALLBACK_ACTION = RAISE
-_VALID_CALLBACK_ACTIONS = (None, IGNORE, PRINT, RAISE, WARN)
+class ExecutionCallback(enum.Enum):
+ """Valid callback actions.
+
+ These can be passed to `seterr` or `errstate` to create callbacks when
+ specific events occur (e.g. an operation produces `NaN`s).
+
+ IGNORE: take no action.
+ PRINT: print a warning to `stdout`.
+ RAISE: raise an error (e.g. `InfOrNanError`).
+ WARN: print a warning using `tf.logging.warn`.
+ """
+
+ IGNORE = "ignore"
+ PRINT = "print"
+ RAISE = "raise"
+ WARN = "warn"
+
+_DEFAULT_CALLBACK_ACTION = ExecutionCallback.RAISE
# TODO(cais): Consider moving this exception class to errors_impl.py.
@@ -139,11 +152,8 @@
the output tensor values.
check_nan: (`bool`) Whether this callback should check for `nan` values in
the output tensor values.
- action: (`str`) Action to be taken by the callback when `inf` or `nan`
- values are detected. Possible values {"raise", "warn", "print"}
- `"raise"`: Raise a `InfOrNanError`.
- `"warn"`: Log a warning using `tf.logging.warn`.
- `"print"`: Print a message to `sys.stdout`.
+ action: (`ExecutionCallback`) Action to be taken by the callback when
+ `inf` or `nan` values are detected.
Raises:
InfOrNanError: iff `inf` or `nan` values are seen in any of `outputs` and
@@ -152,6 +162,7 @@
"""
del attrs, inputs # Not used.
+ action = ExecutionCallback(action)
ctx = context.context()
for index, output in enumerate(outputs):
@@ -180,16 +191,16 @@
continue
error = InfOrNanError(op_type, op_name, index, len(outputs), value)
- if action == "print":
+ if action == ExecutionCallback.PRINT:
print("Warning: %s" % str(error))
- elif action == "warn":
+ elif action == ExecutionCallback.WARN:
logging.warn(str(error))
- elif action == "raise":
+ elif action == ExecutionCallback.RAISE:
raise error
else:
raise ValueError(
"Invalid action for inf_nan_callback: %s. Valid actions are: "
- "{print | warn | raise}" % action)
+ "{PRINT | WARN | RAISE}" % action)
def inf_callback(op_type,
@@ -282,7 +293,7 @@
Example:
```python
- tfe.seterr(inf_or_nan="raise")
+ tfe.seterr(inf_or_nan=ExecutionCallback.RAISE)
a = tf.constant(10.0)
b = tf.constant(0.0)
try:
@@ -290,18 +301,14 @@
except Exception as e:
print("Caught Exception: %s" % e)
- tfe.seterr(inf_or_nan="ignore")
+ tfe.seterr(inf_or_nan=ExecutionCallback.IGNORE)
c = a / b # <-- Does NOT raise exception anymore.
```
Args:
- inf_or_nan: Set action for infinity (`inf`) and NaN (`nan`) values.
- Possible values: `{"ignore", "print", "raise", "warn"}`.
- `"ignore"`: take no action when `inf` values appear.
- `"print"`: print a warning to `stdout`.
- `"raise"`: raise an `InfOrNanError`.
- `"warn"`: print a warning using `tf.logging.warn`.
- A value of `None` leads to no change in the action of the condition.
+ inf_or_nan: An `ExecutionCallback` determining the action for infinity
+ (`inf`) and NaN (`nan`) values. A value of `None` leads to no change in
+ the action of the condition.
Returns:
A dictionary of old actions.
@@ -309,12 +316,8 @@
Raises:
ValueError: If the value of any keyword arguments is invalid.
"""
- if inf_or_nan not in _VALID_CALLBACK_ACTIONS:
- raise ValueError(
- "Invalid action value for inf_or_nan: %s. "
- "Valid actions are %s." % (inf_or_nan, _VALID_CALLBACK_ACTIONS))
-
- old_settings = {"inf_or_nan": "ignore"}
+ inf_or_nan = ExecutionCallback(inf_or_nan) if inf_or_nan is not None else None
+ old_settings = {"inf_or_nan": ExecutionCallback.IGNORE}
default_context = context.context()
carryover_callbacks = []
@@ -336,7 +339,7 @@
default_context.clear_post_execution_callbacks()
for callback in carryover_callbacks:
default_context.add_post_execution_callback(callback)
- if inf_or_nan != "ignore":
+ if inf_or_nan != ExecutionCallback.IGNORE:
default_context.add_post_execution_callback(
functools.partial(inf_nan_callback, action=inf_or_nan))
@@ -351,18 +354,14 @@
```
c = tf.log(0.) # -inf
- with errstate(inf_or_nan="raise"):
+ with errstate(inf_or_nan=ExecutionCallback.RAISE):
tf.log(0.) # <-- Raises InfOrNanError.
```
Args:
- inf_or_nan: Set action for infinity (`inf`) and NaN (`nan`) values.
- Possible values: `{IGNORE, PRINT, RAISE, WARN}`.
- `IGNORE`: take no action when `inf` values appear.
- `PRINT`: print a warning to `stdout`.
- `RAISE`: raise an `InfOrNanError`.
- `WARN`: print a warning using `tf.logging.warn`.
- A value of `None` leads to no change in the action of the condition.
+ inf_or_nan: An `ExecutionCallback` determining the action for infinity
+ (`inf`) and NaN (`nan`) values. A value of `None` leads to no change in
+ the action of the condition.
Yields:
None.
diff --git a/tensorflow/python/eager/execution_callbacks_test.py b/tensorflow/python/eager/execution_callbacks_test.py
index 5594ab5..b8b786a 100644
--- a/tensorflow/python/eager/execution_callbacks_test.py
+++ b/tensorflow/python/eager/execution_callbacks_test.py
@@ -24,6 +24,9 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
+RAISE = execution_callbacks.ExecutionCallback.RAISE
+IGNORE = execution_callbacks.ExecutionCallback.IGNORE
+
def log_zero():
"""Computes `log(0.0)`."""
@@ -33,17 +36,17 @@
class ExecutionCallbacksTest(test.TestCase):
def test_errstate_inf_raise(self):
- with execution_callbacks.errstate(inf_or_nan=execution_callbacks.RAISE):
+ with execution_callbacks.errstate(inf_or_nan=RAISE):
with self.assertRaises(execution_callbacks.InfOrNanError):
log_zero()
def test_errstate_inf_ignore(self):
- with execution_callbacks.errstate(inf_or_nan=execution_callbacks.IGNORE):
+ with execution_callbacks.errstate(inf_or_nan=IGNORE):
self.assertEqual(-float("inf"), log_zero().numpy())
def test_errstate_nesting(self):
- with execution_callbacks.errstate(inf_or_nan=execution_callbacks.RAISE):
- with execution_callbacks.errstate(inf_or_nan=execution_callbacks.IGNORE):
+ with execution_callbacks.errstate(inf_or_nan=RAISE):
+ with execution_callbacks.errstate(inf_or_nan=IGNORE):
self.assertEqual(-float("inf"), log_zero().numpy())
with self.assertRaises(execution_callbacks.InfOrNanError):
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 520c85a..6770f1d 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -149,10 +149,9 @@
outputs: the tensors in the graph which will be outputs to the function
attrs: dict mapping names of attributes to their AttrValue values
"""
- operations = [
- op for op in graph.get_operations()
- if op not in set(arg.op for arg in inputs)
- ]
+ input_ops = set(arg.op for arg in inputs)
+ operations = [op for op in graph.get_operations() if op not in input_ops]
+
fn = pywrap_tensorflow.TF_GraphToFunction_wrapper(
graph._c_graph, # pylint: disable=protected-access
compat.as_str(name),
@@ -340,7 +339,7 @@
TypeError: For invalid positional/keyword argument combinations.
"""
if self._arg_keywords is None or self._num_positional_args is None:
- if self._signature:
+ if self._signature is not None:
if kwargs:
raise NotImplementedError(
"Keyword arguments not supported when calling a "
@@ -748,6 +747,19 @@
return ret
+class UnknownArgument(object):
+ """Signifies an argument which is not currently handled."""
+ pass
+
+
+def _encode_arg_for_serialization(arg):
+ """A representation for this argument, for serializing signatures."""
+ if isinstance(arg, ops.Tensor):
+ return tensor_spec.TensorSpec(arg.shape, arg.dtype)
+ else:
+ return UnknownArgument()
+
+
pywrap_tensorflow.RegisterType("Tensor", ops.Tensor)
pywrap_tensorflow.RegisterType("IndexedSlices", ops.IndexedSlices)
@@ -1163,6 +1175,14 @@
autograph=self._autograph,
arg_names=arg_names),
self._function_attributes)
+ if self._input_signature:
+ python_call_signature = self._input_signature
+ else:
+ python_call_signature = tuple(
+ _encode_arg_for_serialization(arg) for arg in args)
+ # Save information about non-Tensor arguments with the concrete
+ # function. Used to serialize PolymorphicFunctions.
+ graph_function._python_call_signature = python_call_signature # pylint: disable=protected-access
self._function_cache[cache_key] = graph_function
return graph_function, args, kwargs
@@ -1203,19 +1223,18 @@
def defun(func=None, input_signature=None, autograph=True):
"""Compiles a Python function into a callable TensorFlow graph.
- `defun` (short for "define function") trace-compiles a Python function
+ `defun` (short for "define function") compiles a Python function
composed of TensorFlow operations into a callable that executes a `tf.Graph`
containing those operations. The callable produced by `defun` contains only
the subgraph of TensorFlow operations that were executed when the Python
function was called with a particular input signature, defined as a list
of the shapes and dtypes of the Python function's Tensor-valued arguments and
- the values of its non-Tensor Python objects. In particular, `defun` is _not_ a
- compiler for arbitrary Python code.
+ the values of its non-Tensor Python objects.
When eager execution is enabled, the ability to create graphs from Python
functions makes it possible to incrementally trade off debugability and
interactivity for performance. Functions compiled with `defun` cannot be
- inspected with `pdb` and `print` statements; however, executing a graph
+ inspected with `pdb`; however, executing a graph
generated by `defun` sometimes takes less time and memory than eagerly
executing the corresponding Python function, since specifying computations as
graphs allows for optimizations like automatic buffer reuse and
@@ -1306,6 +1325,7 @@
outer graph otherwise.
_Input Signatures_
+
By default, `F = tf.contrib.eager.defun(f)` instantiates a separate graph
for every unique sequence of the shapes and dtypes of Tensor arguments and
the values of Python objects it is invoked with. For example, calling
@@ -1364,6 +1384,7 @@
Tensors as arguments and must not take unnamed keyword arguments (**kwargs).
_Tracing_
+
Be aware that because `F` only logs TensorFlow operations, all the other
Python code that `f` executes will only shape the _construction_ of the graphs
that `F` executes: the Python code won't be executed when the graphs
@@ -1389,6 +1410,7 @@
replace the call to `np.random.randn` with `tf.random_normal((5, 5))`.
_Python Side-Effects_
+
A corollary of the previous discussion on tracing is the following: If a
Python function `f` has Python side-effects, then executing `f` multiple times
will not necessarily be semantically equivalent to executing `F =
@@ -1396,7 +1418,8 @@
that `defun` only captures the subgraph of TensorFlow operations that is
constructed when `f` is called in a graph-building context.
- _Python Control Flow_.
+ _Python Control Flow_
+
The structure of many machine learning computations depend upon whether one is
training or validating, and it is common to nest specialized logic under `if
training:` blocks. By mapping each input signature to a unique graph, `defun`
@@ -1425,27 +1448,26 @@
exact_outputs = lossy_matmul(W, x, training=False)
```
- On the other hand, because `defun` generates graphs by tracing and not by
- source code analysis, it fully unrolls Python `for` and `while` loops,
- potentially creating large graphs. If your Python function has native loops
- that run for many iterations, consider replacing them with `tf.while_loop`
- operations.
+ _TensorFlow Control Flow_
- When constructing graphs, `tf.Tensor` objects cannot be used as Python
- `bool` objects. This means, for example, that you should replace code in `f`
- resembling
+ When `autograph` is `True`, data-dependent control flow is allowed as well.
+ Control flow statements that depend on `Tensor` values are staged into
+ corresponding TensorFlow ops. For example, the following code will work as
+ expected:
```python
-
- if tensor < 10:
- true_fn()
- else:
- false_fn()
+ @tf.contrib.eager.defun
+ def dynamic_rnn_loop(cell, seq):
+ state, output = cell.zero_state()
+ for input in seq:
+ state, output = cell(input, state)
+ return output
```
- with `tf.cond(tensor < 10, true_fn, false_fn)`.
+ For more information see `tf.autograph`.
_Variables_
+
TensorFlow operations related to variable creation and initialization are
automatically lifted out of the graphs generated by `defun`. In practice, this
implies that variable creation and initialization only happen the first time
@@ -1618,12 +1640,19 @@
assert hasattr(original_function, "python_function")
def bound_method_wrapper(*args, **kwargs):
+ """Wraps either a dummy MethodType or a converted AutoGraph function."""
# __wrapped__ allows AutoGraph to swap in a converted function.
wrapped_fn = bound_method_wrapper.__wrapped__
- # If __wrapped__ was not replaced, then call original_function.
- # TODO(b/119246461): This needs to be simplified.
- if tf_inspect.ismethod(wrapped_fn):
+
+ if wrapped_fn is bound_method_wrapper.__original_wrapped__:
+ # If __wrapped__ was not replaced, then call original_function.
wrapped_fn = original_function.python_function
+ if tf_inspect.ismethod(wrapped_fn):
+ wrapped_fn = six.get_unbound_function(wrapped_fn)
+ return wrapped_fn(weak_instance(), *args, **kwargs)
+
+ # If __wrapped__ was replaced, then it is always an unbound function
+ # that takes self as first argument.
return wrapped_fn(weak_instance(), *args, **kwargs)
# pylint: disable=protected-access
diff --git a/tensorflow/python/eager/function_gradients_test.py b/tensorflow/python/eager/function_gradients_test.py
index 9b83f57..98dec0b 100644
--- a/tensorflow/python/eager/function_gradients_test.py
+++ b/tensorflow/python/eager/function_gradients_test.py
@@ -187,7 +187,7 @@
self.assertAllEqual(2, g(constant_op.constant(2.)))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testGraphModeEagerGradError(self):
with context.graph_mode():
def f():
diff --git a/tensorflow/python/eager/function_test.py b/tensorflow/python/eager/function_test.py
index 71afbd2..95777a3 100644
--- a/tensorflow/python/eager/function_test.py
+++ b/tensorflow/python/eager/function_test.py
@@ -428,20 +428,21 @@
self.evaluate(variables.global_variables_initializer())
self.assertEqual(self.evaluate(value), 2.0)
- @test_util.run_in_graph_and_eager_modes
+ @test_util.also_run_as_tf_function
def testInitScopeTensorInitializationInFunction(self):
@def_function.function
def tensor_init():
with ops.init_scope():
const = constant_op.constant(2.0)
+ # Note: this variable bypasses tf.function's variable creation
+ # requirements by bypassing variable_creator_scope by using
+ # ResourceVariable instead of Variable.
self.v = resource_variable_ops.ResourceVariable(const)
return self.v.read_value()
value = tensor_init()
- if not context.executing_eagerly():
- self.evaluate(variables.global_variables_initializer())
- self.assertEqual(self.evaluate(value), 2.0)
+ self.assertAllEqual(value, 2.0)
def testDefunShapeInferenceWithCapturedResourceVariable(self):
v = resource_variable_ops.ResourceVariable([[1, 2], [3, 4]])
@@ -461,6 +462,7 @@
var_t = resource_variable_ops.read_variable_op(var_handle, dtype=v.dtype)
self.assertEqual(var_t.shape, tensor_shape.TensorShape([2, 2]))
+ @test_util.enable_control_flow_v2
def testVariableInLoopInFunction(self):
@function.defun
@@ -543,7 +545,7 @@
self.assertIsInstance(
self.v, resource_variable_ops.ResourceVariable)
- def disabled_testRunMetadata(self):
+ def testRunMetadata(self):
@def_function.function
def f(x):
@@ -578,7 +580,7 @@
return self.v * 2
o = HasAVar()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
call = def_function.function(o.call)
op = call()
self.assertAllEqual(self.evaluate(op), 2.0)
@@ -962,6 +964,7 @@
# construction. Eager's configuration is controlled in `__main__`.
@test_util.run_in_graph_and_eager_modes(
config=config_pb2.ConfigProto(device_count={'CPU': 4}))
+ @test_util.run_v1_only('b/120545219')
def testDeviceAnnotationsRespected(self):
def multi_device_fn():
@@ -1000,6 +1003,7 @@
@test_util.run_in_graph_and_eager_modes(
config=config_pb2.ConfigProto(device_count={'CPU': 2}))
+ @test_util.run_v1_only('b/120545219')
def testCallingGraphFunctionOnDifferentDevice(self):
def func():
diff --git a/tensorflow/python/eager/pywrap_tensor.cc b/tensorflow/python/eager/pywrap_tensor.cc
index 206b96e..30a93fb 100644
--- a/tensorflow/python/eager/pywrap_tensor.cc
+++ b/tensorflow/python/eager/pywrap_tensor.cc
@@ -221,7 +221,7 @@
}
}
tensorflow::Safe_PyObjectPtr value_decrefer;
- if (PyArray_CheckAnyScalarExact(value)) {
+ if (PyArray_IsScalar(value, Generic)) {
// Convert numpy scalars to numpy arrays.
value = PyArray_FromScalar(value, nullptr);
// The returned value needs to be DECREF'd, but the original value was
diff --git a/tensorflow/python/eager/wrap_function_test.py b/tensorflow/python/eager/wrap_function_test.py
index b32b6ca..d34e922 100644
--- a/tensorflow/python/eager/wrap_function_test.py
+++ b/tensorflow/python/eager/wrap_function_test.py
@@ -19,6 +19,7 @@
from tensorflow.python.eager import wrap_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
@@ -70,6 +71,14 @@
f_pruned = f_wrapped.prune(x_in[0], [x_out[0]])
self.assertAllEqual(f_pruned(ops.convert_to_tensor(2.0)), [4.0])
+ def testNoArguments(self):
+
+ def f():
+ return constant_op.constant(1.)
+
+ f_wrapped = wrap_function.wrap_function(f, [])
+ self.assertAllEqual(1.0, f_wrapped())
+
if __name__ == '__main__':
ops.enable_eager_execution()
diff --git a/tensorflow/python/feature_column/feature_column_v2.py b/tensorflow/python/feature_column/feature_column_v2.py
index 6308926..4cc8efa 100644
--- a/tensorflow/python/feature_column/feature_column_v2.py
+++ b/tensorflow/python/feature_column/feature_column_v2.py
@@ -3111,7 +3111,7 @@
'Suggested fix: Use one of sequence_categorical_column_with_*. '
'Given (type {}): {}'.format(self.name, type(self.categorical_column),
self.categorical_column))
- sparse_tensors = self.categorical_column.get_sequence_sparse_tensors(
+ sparse_tensors = self.categorical_column.get_sparse_tensors(
transformation_cache, state_manager)
dense_tensor = self._get_dense_tensor_internal(sparse_tensors,
state_manager)
@@ -3307,7 +3307,7 @@
'Suggested fix A: If you wish to use input_layer, use a '
'non-sequence categorical_column_with_*. '
'Suggested fix B: If you wish to create sequence input, use '
- 'sequence_input_layer instead of input_layer. '
+ 'SequenceFeatureLayer instead of FeatureLayer. '
'Given (type {}): {}'.format(self.name, type(self.categorical_column),
self.categorical_column))
return self._get_dense_tensor_internal(transformation_cache, state_manager)
@@ -3321,12 +3321,12 @@
raise ValueError(
'In embedding_column: {}. '
'categorical_column must be of type SequenceCategoricalColumn '
- 'to use sequence_input_layer. '
+ 'to use SequenceFeatureLayer. '
'Suggested fix: Use one of sequence_categorical_column_with_*. '
'Given (type {}): {}'.format(self.name, type(self.categorical_column),
self.categorical_column))
- dense_tensor = self.get_dense_tensor_internal(transformation_cache,
- state_manager)
+ dense_tensor = self._get_dense_tensor_internal(transformation_cache,
+ state_manager)
sparse_tensors = self.categorical_column.get_sparse_tensors(
transformation_cache, state_manager)
sequence_length = fc_old._sequence_length_from_sparse_tensor( # pylint: disable=protected-access
@@ -4469,8 +4469,8 @@
class SequenceCategoricalColumn(
- FeatureColumn,
- fc_old._CategoricalColumn, # pylint: disable=protected-access
+ CategoricalColumn,
+ fc_old._SequenceCategoricalColumn, # pylint: disable=protected-access
collections.namedtuple('SequenceCategoricalColumn',
('categorical_column'))):
"""Represents sequences of categorical data."""
@@ -4533,7 +4533,7 @@
weight_tensor = sparse_ops.sparse_reshape(weight_tensor, target_shape)
return CategoricalColumn.IdWeightPair(id_tensor, weight_tensor)
- def get_sequence_sparse_tensors(self, transformation_cache, state_manager):
+ def get_sparse_tensors(self, transformation_cache, state_manager):
"""Returns an IdWeightPair.
`IdWeightPair` is a pair of `SparseTensor`s which represents ids and
diff --git a/tensorflow/python/framework/auto_control_deps.py b/tensorflow/python/framework/auto_control_deps.py
index 30dc959..a72ded1 100644
--- a/tensorflow/python/framework/auto_control_deps.py
+++ b/tensorflow/python/framework/auto_control_deps.py
@@ -100,6 +100,7 @@
# graph (but that would mess up devices and collections at least,
# probably other things as well).
self._graph = ops.get_default_graph()
+ self._graph._add_control_dependencies = True # pylint: disable=protected-access
self._n_operations = len(self._graph.get_operations())
return self
@@ -170,6 +171,14 @@
raise RuntimeError(
"Graph changed while trying to add control dependencies.")
+ # pylint: disable=protected-access
+ if hasattr(self._graph, "outer_graph"):
+ outer_val = self._graph.outer_graph._add_control_dependencies
+ self._graph._add_control_dependencies = outer_val
+ else:
+ self._graph._add_control_dependencies = False
+ # pylint: enable=protected-access
+
# map from resource tensor to the last op which used it
last_op_using_resource_tensor = {}
# set of conditional and loop exits
diff --git a/tensorflow/python/framework/auto_control_deps_test.py b/tensorflow/python/framework/auto_control_deps_test.py
index a1dff9e..d81adef 100644
--- a/tensorflow/python/framework/auto_control_deps_test.py
+++ b/tensorflow/python/framework/auto_control_deps_test.py
@@ -24,6 +24,7 @@
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.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import resource_variable_ops
@@ -38,7 +39,7 @@
def testBasic(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
with acd.AutomaticControlDependencies() as c:
v.assign(v + 1)
v.assign(2 * v)
@@ -46,10 +47,11 @@
val = c.mark_as_return(val)
self.assertAllEqual(val.eval(), 4.0)
+ @test_util.run_v1_only("b/120545219")
def testCondMustRun(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
p = array_ops.placeholder(dtype=dtypes.bool)
with acd.AutomaticControlDependencies() as c:
@@ -67,10 +69,11 @@
self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0)
self.assertAllEqual(val.eval(feed_dict={p: True}), 6.0)
+ @test_util.run_v1_only("b/120545219")
def testCondMustRunSeparateRead(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
p = array_ops.placeholder(dtype=dtypes.bool)
with acd.AutomaticControlDependencies() as c:
@@ -90,10 +93,11 @@
one.eval(feed_dict={p: True})
self.assertAllEqual(v.read_value().eval(), 6.0)
+ @test_util.run_v1_only("b/120545219")
def testCondNested(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
p = array_ops.placeholder(dtype=dtypes.bool)
q = array_ops.placeholder(dtype=dtypes.bool)
with acd.AutomaticControlDependencies() as c:
@@ -124,10 +128,11 @@
self.assertAllEqual(val.eval(feed_dict={p: True, q: True}), 7.0)
self.assertAllEqual(val.eval(feed_dict={p: True, q: False}), 8.0)
+ @test_util.run_v1_only("b/120545219")
def testCondOneBranch(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
p = array_ops.placeholder(dtype=dtypes.bool)
with acd.AutomaticControlDependencies() as c:
@@ -144,10 +149,11 @@
self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0)
self.assertAllEqual(val.eval(feed_dict={p: True}), 5.0)
+ @test_util.run_v1_only("b/120545219")
def testCondOneBranchUpdateBefore(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
p = array_ops.placeholder(dtype=dtypes.bool)
with acd.AutomaticControlDependencies() as c:
v.assign(v * 2)
@@ -165,10 +171,11 @@
self.assertAllEqual(val.eval(feed_dict={p: False}), 6.0)
self.assertAllEqual(val.eval(feed_dict={p: True}), 12.0)
+ @test_util.run_v1_only("b/120545219")
def testCondOneBranchUpdateAfter(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
p = array_ops.placeholder(dtype=dtypes.bool)
with acd.AutomaticControlDependencies() as c:
@@ -204,7 +211,7 @@
def testDecorator(self):
with context.graph_mode(), self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
@acd.automatic_control_dependencies
def f():
diff --git a/tensorflow/python/framework/dtypes.py b/tensorflow/python/framework/dtypes.py
index f7a12d2..9a4fe4e 100644
--- a/tensorflow/python/framework/dtypes.py
+++ b/tensorflow/python/framework/dtypes.py
@@ -347,7 +347,7 @@
uint64 = DType(types_pb2.DT_UINT64)
tf_export("dtypes.uint64", "uint64").export_constant(__name__, "uint64")
int16 = DType(types_pb2.DT_INT16)
-tf_export("dtypes.uint16", "int16").export_constant(__name__, "int16")
+tf_export("dtypes.int16", "int16").export_constant(__name__, "int16")
int8 = DType(types_pb2.DT_INT8)
tf_export("dtypes.int8", "int8").export_constant(__name__, "int8")
string = DType(types_pb2.DT_STRING)
diff --git a/tensorflow/python/framework/error_interpolation_test.py b/tensorflow/python/framework/error_interpolation_test.py
index 1b77548..9eaa4a5 100644
--- a/tensorflow/python/framework/error_interpolation_test.py
+++ b/tensorflow/python/framework/error_interpolation_test.py
@@ -23,6 +23,7 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import error_interpolation
from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
from tensorflow.python.framework import traceable_stack
from tensorflow.python.platform import test
from tensorflow.python.util import tf_stack
@@ -112,6 +113,7 @@
self.assertIn("No node-device colocations", summary)
+@test_util.run_v1_only("b/120545219")
class InterpolateFilenamesAndLineNumbersTest(test.TestCase):
def setUp(self):
@@ -193,6 +195,7 @@
self.assertRegexpMatches(interpolated_string, "constant_op.py:[0-9]+.*")
+@test_util.run_v1_only("b/120545219")
class InterpolateDeviceSummaryTest(test.TestCase):
def _fancy_device_function(self, unused_op):
@@ -236,6 +239,7 @@
self.assertRegexpMatches(result, expected_re)
+@test_util.run_v1_only("b/120545219")
class InterpolateColocationSummaryTest(test.TestCase):
def setUp(self):
@@ -260,11 +264,13 @@
self.graph = node_three.graph
+ @test_util.run_v1_only("b/120545219")
def testNodeThreeHasColocationInterpolation(self):
message = "{{colocation_node Three_with_one}}"
result = error_interpolation.interpolate(message, self.graph)
self.assertIn("colocate_with(One)", result)
+ @test_util.run_v1_only("b/120545219")
def testNodeFourHasColocationInterpolationForNodeThreeOnly(self):
message = "{{colocation_node Four_with_three}}"
result = error_interpolation.interpolate(message, self.graph)
@@ -273,12 +279,14 @@
"One", result,
"Node One should not appear in Four_with_three's summary:\n%s" % result)
+ @test_util.run_v1_only("b/120545219")
def testNodeFiveHasColocationInterpolationForNodeOneAndTwo(self):
message = "{{colocation_node Five_with_one_with_two}}"
result = error_interpolation.interpolate(message, self.graph)
self.assertIn("colocate_with(One)", result)
self.assertIn("colocate_with(Two)", result)
+ @test_util.run_v1_only("b/120545219")
def testColocationInterpolationForNodeLackingColocation(self):
message = "{{colocation_node One}}"
result = error_interpolation.interpolate(message, self.graph)
diff --git a/tensorflow/python/framework/func_graph.py b/tensorflow/python/framework/func_graph.py
index f74d072..75a420e 100644
--- a/tensorflow/python/framework/func_graph.py
+++ b/tensorflow/python/framework/func_graph.py
@@ -36,6 +36,7 @@
from tensorflow.python.ops import variable_scope
from tensorflow.python.util import compat
from tensorflow.python.util import nest
+from tensorflow.python.util import tf_contextlib
from tensorflow.python.util import tf_decorator
from tensorflow.python.util.lazy_loader import LazyLoader
@@ -108,36 +109,20 @@
graph = self.outer_graph
- # pylint: disable=protected-access
- # TODO(b/112906995, nareshmodi): distribution strategy depends on inheriting
- # this stack from the default graph even in eager mode. Maybe it should be
- # part of the eager context? This would also allow us to remove a
- # get_default_graph() call from the function cache lookup.
- self._distribution_strategy_stack = list(graph._distribution_strategy_stack)
- # We ignore device placements from any outer scopes while tracing the
- # function when possible, to avoid hard-coding them in the function
- # graph. "Default" placements come from the PartitionedCallOp's placement,
- # so that the same trace of the Python function may be placed on several
- # different devices and saved functions may be placed on new devices when
- # restored.
if context.executing_eagerly():
self.seed = context.global_seed()
- self._xla_compile = (context.context().device_spec.device_type == "TPU")
- if self._distribution_strategy_stack or self._xla_compile:
- self._add_device_to_stack(context.context().device_name)
+ device_type = context.context().device_spec.device_type
+ self._xla_compile = (device_type == "TPU" or device_type == "XLA_GPU"
+ or device_type == "XLA_CPU")
else:
self.seed = graph.seed
self._xla_compile = getattr(graph, "_xla_compile", False)
# TODO(allenl): Figure out if we can remove colocation stack
# specialization (currently used in cond_v2), here and in the cache key.
- self._colocation_stack = graph._colocation_stack.copy()
- if (self._distribution_strategy_stack
- or self._xla_compile
- or device_stack_has_callable(graph._device_function_stack)):
- # Hard-code devices from device functions in the function body
- self._device_function_stack = graph._device_function_stack.copy()
+ self._colocation_stack = graph._colocation_stack.copy() # pylint: disable=protected-access
+
if not self._read_only_collections:
- self._collections = graph._collections
+ self._collections = graph._collections # pylint: disable=protected-access
else:
for collection_name in graph.get_all_collection_keys():
if collection_name not in WHITELIST_COLLECTIONS:
@@ -147,11 +132,63 @@
self._collections[collection_name] = graph.get_collection_ref(
collection_name)
- self._variable_creator_stack = graph._variable_creator_stack
- # Inherit the graph key, since this is used for matching variables in
- # optimizers.
- self._graph_key = graph._graph_key
- # pylint: enable=protected-access
+ def as_default(self):
+ outer_cm = super(FuncGraph, self).as_default()
+
+ @tf_contextlib.contextmanager
+ def inner_cm():
+ """Context manager for copying distribute.Strategy scope information."""
+ graph = ops.get_default_graph()
+ # pylint: disable=protected-access
+ # TODO(b/112906995, nareshmodi): distribution strategy depends on
+ # inheriting this stack from the default graph even in eager mode. Maybe
+ # it should be part of the eager context? This would also allow us to
+ # remove a get_default_graph() call from the function cache lookup.
+ old_strategy_stack = self._distribution_strategy_stack
+ self._distribution_strategy_stack = list(
+ graph._distribution_strategy_stack)
+ # We ignore device placements from any outer scopes while tracing the
+ # function when possible, to avoid hard-coding them in the function
+ # graph. "Default" placements come from the PartitionedCallOp's placement,
+ # so that the same trace of the Python function may be placed on several
+ # different devices and saved functions may be placed on new devices when
+ # restored.
+ old_device_stack = self._device_function_stack
+ if context.executing_eagerly():
+ if self._distribution_strategy_stack or self._xla_compile:
+ self._add_device_to_stack(context.context().device_name)
+ else:
+ if (self._distribution_strategy_stack
+ or self._xla_compile
+ or device_stack_has_callable(graph._device_function_stack)):
+ # Hard-code devices from device functions in the function body
+ self._device_function_stack = graph._device_function_stack.copy()
+
+ old_creator_stack = self._variable_creator_stack
+ self._variable_creator_stack = graph._variable_creator_stack
+ # Inherit the graph key, since this is used for matching variables in
+ # optimizers.
+ old_graph_key = self._graph_key
+ self._graph_key = graph._graph_key
+ # pylint: enable=protected-access
+
+ with outer_cm as g:
+ try:
+ yield g
+ finally:
+ self._distribution_strategy_stack = old_strategy_stack
+ self._device_function_stack = old_device_stack
+ self._variable_creator_stack = old_creator_stack
+ self._graph_key = old_graph_key
+ return inner_cm()
+
+ @property
+ def output_types(self):
+ return [t.dtype for t in self.outputs]
+
+ @property
+ def output_shapes(self):
+ return [t.shape for t in self.outputs]
@property
def variables(self):
diff --git a/tensorflow/python/framework/function.py b/tensorflow/python/framework/function.py
index 622686c..afc11b1 100644
--- a/tensorflow/python/framework/function.py
+++ b/tensorflow/python/framework/function.py
@@ -209,6 +209,7 @@
out_names=None,
shape_func=None,
capture_by_value=False,
+ whitelisted_stateful_ops=None,
**kwargs):
"""Creates _DefinedFunction.
@@ -229,6 +230,8 @@
output shapes.
capture_by_value: Boolean (defaults to False). If True, captured values
will be copied into the function body.
+ whitelisted_stateful_ops: A set of ops that if stateful we ignore and
+ copy into the function body, when `capture_by_value` is True.
**kwargs: The keyword arguments. **kwargs is passed to every call
site of this function.
@@ -244,6 +247,9 @@
self._out_names = out_names
self._shape_func = shape_func
self._capture_by_value = capture_by_value
+ self._whitelisted_stateful_ops = whitelisted_stateful_ops
+ if self._whitelisted_stateful_ops is None:
+ self._whitelisted_stateful_ops = set()
self._extra_kwargs = kwargs
# Constructed only when C API is disabled, lazily
self._definition = None
@@ -340,8 +346,13 @@
return
temp_graph = func_graph_from_py_func(
- self._func, self._arg_names, self._arg_types, self._func_name,
- self._capture_by_value, self._caller_device)
+ self._func,
+ self._arg_names,
+ self._arg_types,
+ self._func_name,
+ self._capture_by_value,
+ self._caller_device,
+ whitelisted_stateful_ops=self._whitelisted_stateful_ops)
self._extra_inputs = temp_graph.extra_inputs
# pylint: disable=protected-access
@@ -625,9 +636,11 @@
function argument and the caller passes in the captured tensor.
"""
- def __init__(self, name, capture_by_value, *args, **kwargs):
+ def __init__(self, name, capture_by_value, whitelisted_stateful_ops, *args,
+ **kwargs):
super(_FuncGraph, self).__init__(*args, **kwargs)
self._capture_by_value = capture_by_value
+ self._whitelisted_stateful_ops = whitelisted_stateful_ops
self._building_function = True
self._outer_graph = ops.get_default_graph()
self._vscope = vs.get_variable_scope()
@@ -785,7 +798,7 @@
# pylint: disable=protected-access
op_def = graph_to_function_def._get_op_def(op)
# pylint: enable=protected-access
- if op_def.is_stateful:
+ if op_def.is_stateful and op not in self._whitelisted_stateful_ops:
raise ValueError("Cannot capture a stateful node (name:%s, type:%s) "
"by value." % (op.name, op.type))
elif op.type in ("Placeholder", "PlaceholderV2"):
@@ -807,10 +820,17 @@
return captured_op
-def func_graph_from_py_func(func, arg_names, arg_types, name=None,
- capture_by_value=False, device=None,
- colocation_stack=None, container=None,
- collections_ref=None, arg_shapes=None):
+def func_graph_from_py_func(func,
+ arg_names,
+ arg_types,
+ name=None,
+ capture_by_value=False,
+ device=None,
+ colocation_stack=None,
+ container=None,
+ collections_ref=None,
+ arg_shapes=None,
+ whitelisted_stateful_ops=None):
"""Returns a _FuncGraph generated from `func`.
Args:
@@ -828,6 +848,8 @@
collections_ref: A reference to a collections dict the _FuncGraph should
use internally.
arg_shapes: A sequence of the function's argument shapes.
+ whitelisted_stateful_ops: A set of ops that if stateful we ignore and
+ re-create.
Returns:
A _FuncGraph.
@@ -837,7 +859,7 @@
"""
if not name:
name = function_utils.get_func_name(func)
- func_graph = _FuncGraph(name, capture_by_value)
+ func_graph = _FuncGraph(name, capture_by_value, whitelisted_stateful_ops)
with func_graph.as_default(), ops.device(device):
# pylint: disable=protected-access
@@ -971,17 +993,18 @@
name = kwargs.pop("name", None)
g = ops.get_default_graph()
func_name = sig.name
+ if name is None:
+ name = func_name
attrs = _parse_kwargs_as_attrs(func_name, **kwargs)
output_types = [dtypes.DType(x.type) for x in sig.output_arg]
- with ops.name_scope(name, func_name, inputs) as name:
- op = g.create_op(
- func_name,
- list(inputs),
- output_types,
- name=name,
- attrs=attrs,
- op_def=sig,
- compute_shapes=False)
+ op = g.create_op(
+ func_name,
+ list(inputs),
+ output_types,
+ name=name,
+ attrs=attrs,
+ op_def=sig,
+ compute_shapes=False)
if op.outputs:
if len(op.outputs) == 1:
ret = op.outputs[0]
@@ -1024,12 +1047,13 @@
c_func = c_api.TF_FunctionImportFunctionDef(serialized)
result._c_func = c_api_util.ScopedTFFunction(c_func)
result._extra_inputs = []
+ result._op_def = fdef.signature
# pylint: enable=protected-access
return result
-def _from_library(lib):
+def from_library(lib):
"""Creates _DefinedFunctions initialized from a FunctionDefLibrary proto.
This method handles assigning the correct gradient functions to each
diff --git a/tensorflow/python/framework/function_test.py b/tensorflow/python/framework/function_test.py
index d71f06e..7543376 100644
--- a/tensorflow/python/framework/function_test.py
+++ b/tensorflow/python/framework/function_test.py
@@ -1054,6 +1054,28 @@
self.assertFalse(all(val3 == val1))
self.assertFalse(all(val4 == val2))
+ def testStatefulFunctionWithWhitelisting(self):
+ t = random_ops.random_uniform([100], maxval=10, dtype=dtypes.int32)
+
+ @function.Defun(capture_by_value=True)
+ def StatefulFn():
+ return t + constant_op.constant(3, dtype=dtypes.int32)
+
+ # First time we try to capture a stateful RandomUniform op.
+ with self.assertRaisesRegexp(ValueError, "Cannot capture a stateful node"):
+ res = StatefulFn()
+
+ # This time we whitelist this op, so that its recreated.
+ @function.Defun(capture_by_value=True, whitelisted_stateful_ops=set([t.op]))
+ def StatefulFn2():
+ return t + constant_op.constant(3, dtype=dtypes.int32)
+
+ res = StatefulFn2()
+ with session.Session() as sess:
+ r = sess.run(res)
+ for i in r:
+ self.assertGreaterEqual(i, 3)
+
@test_util.run_deprecated_v1
def testSameFunctionOnTwoDevices(self):
@@ -1265,7 +1287,7 @@
gradients_impl.gradients([f1, f2, f3, f4], c)
library = g.as_graph_def().library
- new_funcs = function._from_library(library)
+ new_funcs = function.from_library(library)
def CheckNewFunc(func):
new_func = [f for f in new_funcs if f.name == func.name]
@@ -1281,7 +1303,7 @@
def testFromLibraryEmptyLib(self):
library = function_pb2.FunctionDefLibrary()
- self.assertEqual(len(function._from_library(library)), 0)
+ self.assertEqual(len(function.from_library(library)), 0)
def testFromLibraryMissingFuncDef(self):
@@ -1305,7 +1327,7 @@
with self.assertRaisesRegexp(
ValueError,
"FunctionDefLibrary missing 'G1_[0-9a-zA-Z]{8,11}' FunctionDef"):
- function._from_library(library)
+ function.from_library(library)
# Create invalid function def that is missing F1 function def
library = function_pb2.FunctionDefLibrary()
@@ -1315,7 +1337,7 @@
with self.assertRaisesRegexp(
ValueError,
"FunctionDefLibrary missing 'F1_[0-9a-zA-Z]{8,11}' FunctionDef"):
- function._from_library(library)
+ function.from_library(library)
def testFromLibraryCyclicGradFuncs(self):
@@ -1344,7 +1366,7 @@
with self.assertRaisesRegexp(
ValueError, "FunctionDefLibrary contains cyclic gradient functions!"):
- function._from_library(library)
+ function.from_library(library)
def testExperimentalAttrs(self):
diff --git a/tensorflow/python/framework/graph_util_test.py b/tensorflow/python/framework/graph_util_test.py
index 4e7408a..dd26b8a 100644
--- a/tensorflow/python/framework/graph_util_test.py
+++ b/tensorflow/python/framework/graph_util_test.py
@@ -103,7 +103,7 @@
self.assertDeviceEqual(var_5.device, "/device:GPU:0")
self.assertDeviceEqual(var_6.device, "/device:CPU:0")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNestedDeviceFunctions(self):
with ops.Graph().as_default():
var_0 = variables.VariableV1(0)
diff --git a/tensorflow/python/framework/importer.py b/tensorflow/python/framework/importer.py
index 98c7aec..c737bd4 100644
--- a/tensorflow/python/framework/importer.py
+++ b/tensorflow/python/framework/importer.py
@@ -442,11 +442,9 @@
_ProcessNewOps(graph)
if graph_def.library and graph_def.library.function:
- # pylint: disable=protected-access
- functions = function._from_library(graph_def.library)
+ functions = function.from_library(graph_def.library)
for f in functions:
f.add_to_graph(graph)
- # pylint: enable=protected-access
# Treat input mappings that don't appear in the graph as an error, because
# they are likely to be due to a typo.
diff --git a/tensorflow/python/framework/meta_graph_test.py b/tensorflow/python/framework/meta_graph_test.py
index 46ce461..e6e8788 100644
--- a/tensorflow/python/framework/meta_graph_test.py
+++ b/tensorflow/python/framework/meta_graph_test.py
@@ -528,7 +528,7 @@
actual_grad_value = self.evaluate(grad)
self.assertEqual(expected_grad_value, actual_grad_value)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testImportWhileLoopInWhileLoop(self):
# Create a simple while loop.
with ops.Graph().as_default():
diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index fa30693..27c56ef 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -2896,6 +2896,9 @@
self._last_loss_reduction = None
self._container = ""
self._registered_ops = op_def_registry.get_registered_ops()
+ # Set to True if this graph is being built in an
+ # AutomaticControlDependencies context.
+ self._add_control_dependencies = False
# TODO(skyewm): fold as much of the above as possible into the C
# implementation
diff --git a/tensorflow/python/framework/ops_test.py b/tensorflow/python/framework/ops_test.py
index 7baa02b..2d7ee1a 100644
--- a/tensorflow/python/framework/ops_test.py
+++ b/tensorflow/python/framework/ops_test.py
@@ -605,6 +605,7 @@
x.op._update_input(1, x) # pylint: disable=protected-access
@test_util.enable_control_flow_v2
+ @test_util.run_v1_only("b/120545219")
def testAddWhileInput(self):
@eager_function.defun
def test():
@@ -614,6 +615,9 @@
self.assertEqual(while_op.type, "While")
orig_num_inputs = len(while_op.inputs)
+ # Make sure we can handle the while op having a control input.
+ while_op._add_control_input(constant_op.constant(0).op)
+
new_input1 = constant_op.constant(1.0)
new_input2 = constant_op.constant(True)
@@ -780,7 +784,7 @@
self.assertEqual(op3.name, "myop_2")
self.assertEqual(op4.name, "myop_1_1")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCond(self):
g = ops.Graph()
with g.as_default():
@@ -810,7 +814,7 @@
"cond/cond_text")
# pylint: enable=protected-access
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileLoop(self):
g = ops.Graph()
with g.as_default():
@@ -840,7 +844,7 @@
"myloop/while_context")
# pylint: enable=protected-access
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileLoopWithInternalControlDep(self):
g = ops.Graph()
with g.as_default():
@@ -864,7 +868,7 @@
# Internal control dep is preserved
self.assertEqual(op.control_inputs, [c])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileLoopWithExternalControlDep(self):
g = ops.Graph()
with g.as_default():
@@ -2283,7 +2287,7 @@
self.assertEqual(4, int(compiled_outer(inner=compiled_inner)))
self.assertEqual(7, int(compiled_outer(inner=compiled_inner)))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testFallsBackToGlobalGraphWhenAllGraphsAreBuildingFunctions(self):
with context.graph_mode():
ops.reset_default_graph()
@@ -2994,7 +2998,7 @@
class EnableEagerExecutionTest(test_util.TensorFlowTestCase):
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBadArgumentsToEnableEagerExecution(self):
with self.assertRaisesRegexp(TypeError, "config must be a tf.ConfigProto"):
ops.enable_eager_execution(context.DEVICE_PLACEMENT_SILENT)
diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index d91f7b0..d460168 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -634,7 +634,9 @@
bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
const string& parameters, const std::vector<string>& output_sizes,
const string& eager_not_allowed_error) {
- strings::StrAppend(&result_, "@_dispatch.add_dispatch_list\n");
+ if (api_def_.visibility() == ApiDef::VISIBLE) {
+ strings::StrAppend(&result_, "@_dispatch.add_dispatch_list\n");
+ }
AddExport();
AddDefLine(function_name_, parameters);
AddDocStringDescription();
diff --git a/tensorflow/python/framework/sparse_tensor.py b/tensorflow/python/framework/sparse_tensor.py
index 3643fc5..5e1a95a 100644
--- a/tensorflow/python/framework/sparse_tensor.py
+++ b/tensorflow/python/framework/sparse_tensor.py
@@ -244,7 +244,7 @@
SparseTensorValue = collections.namedtuple(
"SparseTensorValue", ["indices", "values", "dense_shape"])
-tf_export("SparseTensorValue")(SparseTensorValue)
+tf_export(v1=["SparseTensorValue"])(SparseTensorValue)
pywrap_tensorflow.RegisterType("SparseTensorValue", SparseTensorValue)
diff --git a/tensorflow/python/framework/subscribe_test.py b/tensorflow/python/framework/subscribe_test.py
index 61c6ea6..a74e96f 100644
--- a/tensorflow/python/framework/subscribe_test.py
+++ b/tensorflow/python/framework/subscribe_test.py
@@ -215,7 +215,7 @@
self.assertIn('graph2', shared)
self.assertIn('graph3', shared)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testSubscribeVariable(self):
"""Confirm that variables can be subscribed."""
v1 = variables.VariableV1(0.0)
@@ -254,7 +254,7 @@
# Make sure the values read from the variable match the expected ones.
self.assertEqual([0.0, 3.0], shared)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testResourceType(self):
"""Confirm that subscribe correctly handles tensors with 'resource' type."""
tensor_array = tensor_array_ops.TensorArray(
@@ -344,7 +344,7 @@
self.assertEqual(add.device, add_sub.device)
self.assertEqual(mul.device, mul_sub.device)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_subscribe_tensors_within_control_flow_context(self):
"""Side effect ops are added with the same control flow context."""
c1 = constant_op.constant(10)
diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index fc1a5fb..d06e1f5 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -54,6 +54,7 @@
from tensorflow.python.client import device_lib
from tensorflow.python.client import session
from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
from tensorflow.python.eager import tape
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import dtypes
@@ -66,8 +67,8 @@
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import versions
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import script_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import googletest
from tensorflow.python.platform import tf_logging as logging
@@ -76,6 +77,7 @@
from tensorflow.python.util import deprecation
from tensorflow.python.util import memory
from tensorflow.python.util import nest
+from tensorflow.python.util import tf_decorator
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.protobuf import compare
from tensorflow.python.util.tf_export import tf_export
@@ -406,42 +408,12 @@
"""
def wrapper(*args, **kwargs):
- enable_cond_v2_old = control_flow_ops.ENABLE_COND_V2
- enable_while_v2_old = control_flow_ops.ENABLE_WHILE_V2
- enable_tensor_array_v2_old = tensor_array_ops.ENABLE_TENSOR_ARRAY_V2
- control_flow_ops.ENABLE_COND_V2 = True
- control_flow_ops.ENABLE_WHILE_V2 = True
- tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 = True
+ enable_control_flow_v2_old = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
try:
fn(*args, **kwargs)
finally:
- control_flow_ops.ENABLE_COND_V2 = enable_cond_v2_old
- control_flow_ops.ENABLE_WHILE_V2 = enable_while_v2_old
- tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 = enable_tensor_array_v2_old
-
- return wrapper
-
-
-def enable_tensor_array_v2(fn):
- """Decorator for enabling _GraphTensorArrayV2 on a test.
-
- Note this enables _GraphTensorArrayV2 after running the test class's
- setup/teardown methods.
-
- Args:
- fn: the function to be wrapped
-
- Returns:
- The wrapped function
- """
-
- def wrapper(*args, **kwargs):
- enable_tensor_array_v2_old = tensor_array_ops.ENABLE_TENSOR_ARRAY_V2
- tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 = True
- try:
- fn(*args, **kwargs)
- finally:
- tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 = enable_tensor_array_v2_old
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = enable_control_flow_v2_old
return wrapper
@@ -490,11 +462,12 @@
Returns:
cls with new test methods added
"""
- if control_flow_ops.ENABLE_WHILE_V2 and control_flow_ops.ENABLE_COND_V2:
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2:
return cls
for name, value in cls.__dict__.copy().items():
- if (callable(value) and name.startswith("test") and
+ if (callable(value) and
+ name.startswith(unittest.TestLoader.testMethodPrefix) and
not getattr(value, "_disable_control_flow_v2", False)):
setattr(cls, name + "WithControlFlowV2", enable_control_flow_v2(value))
return cls
@@ -893,8 +866,10 @@
"""Execute all test methods in the given class with and without eager."""
base_decorator = run_in_graph_and_eager_modes
for name, value in cls.__dict__.copy().items():
- if callable(value) and name.startswith("test") and not (
- name.startswith("testSkipEager") or name.startswith("test_skip_eager")):
+ if (callable(value) and
+ name.startswith(unittest.TestLoader.testMethodPrefix) and
+ not (name.startswith("testSkipEager")
+ or name.startswith("test_skip_eager"))):
setattr(cls, name, base_decorator(value))
return cls
@@ -1006,6 +981,58 @@
return decorator
+def py_func_if_in_function(f):
+
+ def decorated(*args, **kwds):
+ if not ops.get_default_graph()._building_function:
+ return f(*args, **kwds)
+
+ tensor_args, tensor_indices = zip(
+ *[(x, i) for i, x in enumerate(args)
+ if isinstance(x, (ops.Tensor, variables.Variable))])
+
+ def inner_f(*inner_tensor_args):
+ my_args = list(args)
+ for i, n in zip(tensor_indices, inner_tensor_args):
+ my_args[i] = n
+ return f(*my_args, **kwds)
+
+ return script_ops.py_func(inner_f, tensor_args, [])
+
+ return tf_decorator.make_decorator(f, decorated)
+
+
+def also_run_as_tf_function(f):
+ """Runs the decorated test twice--once as is, once inside a tf.function.
+
+ This allows you to run a test both in eager execution and inside a
+ tf.function, exercising the two execution modes supported in tf 2.0. The test
+ assertions are automatically done inside tf.py_funcs, and tf.function ensures
+ that they run in the proper order and with the proper side effects.
+
+ Currently variable creation is not supported in tests annotated with this
+ decorator since it's tricky to ensure the variable doesn't get repeatedly
+ created when retracing the tf.function.
+
+ Args:
+ f: the test method to be decorated
+
+ Returns:
+ The decorated test method, which will run both in eager and inside a
+ tf.function.
+ """
+
+ def decorated(*args, **kwds):
+ with context.eager_mode():
+ # Running in eager mode
+ f(*args, **kwds)
+
+ defun_f = def_function.function(f)
+ defun_f(*args, **kwds)
+
+ return decorated
+
+
def run_deprecated_v1(func=None):
"""Execute the decorated test in graph mode.
@@ -1059,7 +1086,16 @@
def decorator(f):
if tf_inspect.isclass(f):
- raise ValueError("`run_v1_only` only supports test methods.")
+ setup = f.__dict__.get("setUp")
+ if setup is not None:
+ setattr(f, "setUp", decorator(setup))
+
+ for name, value in f.__dict__.copy().items():
+ if (callable(value) and
+ name.startswith(unittest.TestLoader.testMethodPrefix)):
+ setattr(f, name, decorator(value))
+
+ return f
def decorated(self, *args, **kwargs):
if tf2.enabled():
@@ -1265,6 +1301,63 @@
return output_data
+class FakeEagerSession(object):
+ """Fake session so tests that conditionally use placeholders can use eager.
+
+ There are a number of tests that conditionally use placeholders for shape
+ inference. The pattern is demonstrated here:
+
+ ```python
+ with self.cached_session() as sess:
+ if static_shape:
+ y = math_ops.matmul(x, ...)
+ feed_dict = {}
+ else:
+ x_ph = array_ops.placeholder(...)
+ y = math_ops.matmul(x_ph, ...)
+ feed_dict = {x_ph: x}
+ val = sess.run(y, feed_dict=feed_dict)
+ ```
+
+ Since the feed_dict is empty when not using placeholders we should be able to
+ call self.evaluate(), however this requires rewriting the test case.
+ This class shold be considered a stop-gap solution to get tests running with
+ eager with minimal changes to the actual test.
+ """
+
+ def __init__(self, test_case):
+ self._test_case = test_case
+
+ def run(self, fetches, *args, **kwargs):
+ """Evalaute `fetches`.
+
+ Fail if additional args are specified.
+
+ Args:
+ fetches: A Tensor or a nested list/tuple of Tensors.
+ *args: Positional arguments
+ **kwargs: Keyword arguments
+
+ Raises:
+ RuntimeError: If args or kwargs are specified.
+
+ Returns:
+ Tensors as numpy values.
+ """
+ feed_dict = kwargs.pop("feed_dict", {})
+ if feed_dict:
+ raise RuntimeError(
+ "feed_dict is not supported when eager execution is enabled "
+ "(in this case, sess.run(t) is shorthand for t.numpy()")
+
+ if args or kwargs:
+ raise RuntimeError(
+ "Optional args are not supported when eager execution is enabled "
+ "(in this case, sess.run(t) is shorthand for t.numpy()")
+
+ return self._test_case.evaluate(fetches)
+
+
class ErrorLoggingSession(session.Session):
"""Wrapper around a Session that logs errors in run().
"""
@@ -1306,6 +1399,10 @@
ops.reset_default_graph()
random_seed.set_random_seed(random_seed.DEFAULT_GRAPH_SEED)
+ # Avoiding calling setUp() for the poorly named test_session method.
+ if self.id().endswith(".test_session"):
+ self.skipTest("Not a test.")
+
def tearDown(self):
for thread in self._threads:
thread.check_termination()
@@ -1572,7 +1669,7 @@
the graph building and execution code in a test case.
"""
if context.executing_eagerly():
- yield None
+ yield FakeEagerSession(self)
else:
sess = self._get_cached_session(
graph, config, force_gpu, crash_if_inconsistent_args=True)
@@ -1591,7 +1688,6 @@
"""Use cached_session instead."""
if self.id().endswith(".test_session"):
self.skipTest("Not a test.")
-
if context.executing_eagerly():
yield None
else:
@@ -1714,8 +1810,8 @@
return ret
-# pylint: enable=invalid-name
-
+ # pylint: enable=invalid-name
+ @py_func_if_in_function
def assertNear(self, f1, f2, err, msg=None):
"""Asserts that two floats are near each other.
@@ -1734,6 +1830,7 @@
"%f != %f +/- %f%s" % (f1, f2, err, " (%s)" % msg
if msg is not None else ""))
+ @py_func_if_in_function
def assertArrayNear(self, farray1, farray2, err, msg=None):
"""Asserts that two float arrays are near each other.
@@ -1753,6 +1850,7 @@
def _NDArrayNear(self, ndarray1, ndarray2, err):
return np.linalg.norm(ndarray1 - ndarray2) < err
+ @py_func_if_in_function
def assertNDArrayNear(self, ndarray1, ndarray2, err, msg=None):
"""Asserts that two numpy arrays have near values.
@@ -1768,7 +1866,7 @@
# If a is a tensor then convert it to ndarray
if isinstance(a, ops.Tensor):
if isinstance(a, ops._EagerTensorBase):
- return a.numpy()
+ a = a.numpy()
else:
a = self.evaluate(a)
if not isinstance(a, np.ndarray):
@@ -1890,6 +1988,7 @@
e.args = ((e.args[0] + " : " + msg,) + e.args[1:])
raise
+ @py_func_if_in_function
def assertAllClose(self, a, b, rtol=1e-6, atol=1e-6, msg=None):
"""Asserts that two structures of numpy arrays or Tensors, have near values.
@@ -1915,6 +2014,7 @@
"""
self._assertAllCloseRecursive(a, b, rtol=rtol, atol=atol, msg=msg)
+ @py_func_if_in_function
def assertAllCloseAccordingToType(self,
a,
b,
@@ -1962,6 +2062,7 @@
self.assertAllClose(a, b, rtol=rtol, atol=atol, msg=msg)
+ @py_func_if_in_function
def assertNotAllClose(self, a, b, **kwargs):
"""Assert that two numpy arrays, or or Tensors, do not have near values.
@@ -1980,6 +2081,7 @@
return
raise AssertionError("The two values are close at all elements")
+ @py_func_if_in_function
def assertAllEqual(self, a, b, msg=None):
"""Asserts that two numpy arrays or Tensors have the same values.
@@ -2022,6 +2124,7 @@
msgs.append("not equal rhs = {}".format(y))
np.testing.assert_array_equal(a, b, err_msg="\n".join(msgs))
+ @py_func_if_in_function
def assertAllGreater(self, a, comparison_target):
"""Assert element values are all greater than a target value.
@@ -2033,6 +2136,7 @@
a = self._GetNdArray(a)
self.assertGreater(np.min(a), comparison_target)
+ @py_func_if_in_function
def assertAllLess(self, a, comparison_target):
"""Assert element values are all less than a target value.
@@ -2044,6 +2148,7 @@
a = self._GetNdArray(a)
self.assertLess(np.max(a), comparison_target)
+ @py_func_if_in_function
def assertAllGreaterEqual(self, a, comparison_target):
"""Assert element values are all greater than or equal to a target value.
@@ -2055,6 +2160,7 @@
a = self._GetNdArray(a)
self.assertGreaterEqual(np.min(a), comparison_target)
+ @py_func_if_in_function
def assertAllLessEqual(self, a, comparison_target):
"""Assert element values are all less than or equal to a target value.
@@ -2097,6 +2203,7 @@
lines.append(prefix + "...")
return lines
+ @py_func_if_in_function
def assertAllInRange(self,
target,
lower_bound,
@@ -2155,6 +2262,7 @@
"Subscript(s) and value(s) of the offending elements:\n" +
"\n".join(self._format_subscripts(violation_subscripts, target)))
+ @py_func_if_in_function
def assertAllInSet(self, target, expected_set):
"""Assert that elements of a Tensor are all in a given closed set.
@@ -2176,6 +2284,7 @@
raise AssertionError("%d unique element(s) are not in the set %s: %s" %
(np.size(diff), expected_set, diff))
+ @py_func_if_in_function
def assertDTypeEqual(self, target, expected_dtype):
"""Assert ndarray data type is equal to expected.
diff --git a/tensorflow/python/grappler/item_test.py b/tensorflow/python/grappler/item_test.py
index 78604b2..c02fd9f 100644
--- a/tensorflow/python/grappler/item_test.py
+++ b/tensorflow/python/grappler/item_test.py
@@ -108,7 +108,7 @@
newest_tf_item = grappler_item.tf_item
self.assertEqual(new_tf_item, newest_tf_item)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testColocationContraints(self):
with ops.Graph().as_default() as g:
c = constant_op.constant([10])
diff --git a/tensorflow/python/grappler/memory_optimizer_test.py b/tensorflow/python/grappler/memory_optimizer_test.py
index 6eb16fb..e2864eb 100644
--- a/tensorflow/python/grappler/memory_optimizer_test.py
+++ b/tensorflow/python/grappler/memory_optimizer_test.py
@@ -62,7 +62,7 @@
self.assertEqual(len(graph.node), graph_size)
self.assertItemsEqual([node.name for node in graph.node], nodes)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testSimpleSwap(self):
"""Check that the swap annotations are followed."""
a = variables.VariableV1(10, name='a')
diff --git a/tensorflow/python/grappler/tf_optimizer_test.py b/tensorflow/python/grappler/tf_optimizer_test.py
index 06ccaa8..8186c81 100644
--- a/tensorflow/python/grappler/tf_optimizer_test.py
+++ b/tensorflow/python/grappler/tf_optimizer_test.py
@@ -57,7 +57,7 @@
self.assertEqual(len(graph.node), 1)
self.assertItemsEqual([node.name for node in graph.node], ['d'])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testKeepNodes(self):
g = ops.Graph()
with g.as_default():
@@ -86,7 +86,7 @@
self.assertEqual(len(optimized_graph_nodes), len(expected_nodes))
self.assertAllInSet(optimized_graph_nodes, expected_nodes)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testLoops(self):
g = ops.Graph()
with g.as_default():
diff --git a/tensorflow/python/keras/BUILD b/tensorflow/python/keras/BUILD
index 602cf79..c056996 100755
--- a/tensorflow/python/keras/BUILD
+++ b/tensorflow/python/keras/BUILD
@@ -301,15 +301,23 @@
)
py_test(
+ name = "metrics_functional_test",
+ size = "medium",
+ srcs = ["metrics_functional_test.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":keras",
+ "//tensorflow/python:client_testlib",
+ "//third_party/py/numpy",
+ ],
+)
+
+py_test(
name = "metrics_test",
size = "medium",
srcs = ["metrics_test.py"],
+ shard_count = 4,
srcs_version = "PY2AND3",
- tags = [
- "manual",
- "no_oss",
- "notap",
- ],
deps = [
":keras",
"//tensorflow/python:client_testlib",
@@ -361,7 +369,7 @@
name = "convolutional_test",
size = "large",
srcs = ["layers/convolutional_test.py"],
- shard_count = 4,
+ shard_count = 11,
srcs_version = "PY2AND3",
deps = [
":keras",
@@ -401,7 +409,7 @@
name = "core_test",
size = "medium",
srcs = ["layers/core_test.py"],
- shard_count = 2,
+ shard_count = 3,
srcs_version = "PY2AND3",
deps = [
":keras",
@@ -466,6 +474,7 @@
name = "normalization_test",
size = "medium",
srcs = ["layers/normalization_test.py"],
+ shard_count = 3,
srcs_version = "PY2AND3",
tags = ["notsan"],
deps = [
@@ -524,8 +533,9 @@
py_test(
name = "recurrent_test",
- size = "large",
+ size = "medium",
srcs = ["layers/recurrent_test.py"],
+ shard_count = 4,
srcs_version = "PY2AND3",
deps = [
":keras",
diff --git a/tensorflow/python/keras/activations_test.py b/tensorflow/python/keras/activations_test.py
index 6b7bfb6..33001f4 100644
--- a/tensorflow/python/keras/activations_test.py
+++ b/tensorflow/python/keras/activations_test.py
@@ -31,6 +31,7 @@
return e / np.sum(e)
+@test_util.run_all_in_graph_and_eager_modes
class KerasActivationsTest(test.TestCase):
def test_serialization(self):
@@ -46,12 +47,11 @@
assert fn == ref_fn
def test_softmax(self):
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.softmax(x)])
- test_values = np.random.random((2, 5))
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.softmax(x)])
+ test_values = np.random.random((2, 5))
- result = f([test_values])[0]
+ result = f([test_values])[0]
expected = _ref_softmax(test_values[0])
self.assertAllClose(result[0], expected, rtol=1e-05)
@@ -60,40 +60,36 @@
keras.activations.softmax(x)
def test_temporal_softmax(self):
- with self.cached_session():
- x = keras.backend.placeholder(shape=(2, 2, 3))
- f = keras.backend.function([x], [keras.activations.softmax(x)])
- test_values = np.random.random((2, 2, 3)) * 10
- result = f([test_values])[0]
+ x = keras.backend.placeholder(shape=(2, 2, 3))
+ f = keras.backend.function([x], [keras.activations.softmax(x)])
+ test_values = np.random.random((2, 2, 3)) * 10
+ result = f([test_values])[0]
expected = _ref_softmax(test_values[0, 0])
self.assertAllClose(result[0, 0], expected, rtol=1e-05)
- @test_util.run_deprecated_v1
def test_selu(self):
x = keras.backend.placeholder(ndim=2)
f = keras.backend.function([x], [keras.activations.selu(x)])
alpha = 1.6732632423543772848170429916717
scale = 1.0507009873554804934193349852946
- with self.cached_session():
- positive_values = np.array([[1, 2]], dtype=keras.backend.floatx())
- result = f([positive_values])[0]
- self.assertAllClose(result, positive_values * scale, rtol=1e-05)
+ positive_values = np.array([[1, 2]], dtype=keras.backend.floatx())
+ result = f([positive_values])[0]
+ self.assertAllClose(result, positive_values * scale, rtol=1e-05)
- negative_values = np.array([[-1, -2]], dtype=keras.backend.floatx())
- result = f([negative_values])[0]
- true_result = (np.exp(negative_values) - 1) * scale * alpha
- self.assertAllClose(result, true_result)
+ negative_values = np.array([[-1, -2]], dtype=keras.backend.floatx())
+ result = f([negative_values])[0]
+ true_result = (np.exp(negative_values) - 1) * scale * alpha
+ self.assertAllClose(result, true_result)
def test_softplus(self):
def softplus(x):
return np.log(np.ones_like(x) + np.exp(x))
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.softplus(x)])
- test_values = np.random.random((2, 5))
- result = f([test_values])[0]
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.softplus(x)])
+ test_values = np.random.random((2, 5))
+ result = f([test_values])[0]
expected = softplus(test_values)
self.assertAllClose(result, expected, rtol=1e-05)
@@ -101,11 +97,10 @@
def softsign(x):
return np.divide(x, np.ones_like(x) + np.absolute(x))
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.softsign(x)])
- test_values = np.random.random((2, 5))
- result = f([test_values])[0]
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.softsign(x)])
+ test_values = np.random.random((2, 5))
+ result = f([test_values])[0]
expected = softsign(test_values)
self.assertAllClose(result, expected, rtol=1e-05)
@@ -118,68 +113,60 @@
return z / (1 + z)
sigmoid = np.vectorize(ref_sigmoid)
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.sigmoid(x)])
- test_values = np.random.random((2, 5))
- result = f([test_values])[0]
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.sigmoid(x)])
+ test_values = np.random.random((2, 5))
+ result = f([test_values])[0]
expected = sigmoid(test_values)
self.assertAllClose(result, expected, rtol=1e-05)
- @test_util.run_deprecated_v1
def test_hard_sigmoid(self):
def ref_hard_sigmoid(x):
x = (x * 0.2) + 0.5
z = 0.0 if x <= 0 else (1.0 if x >= 1 else x)
return z
hard_sigmoid = np.vectorize(ref_hard_sigmoid)
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.hard_sigmoid(x)])
- test_values = np.random.random((2, 5))
- result = f([test_values])[0]
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.hard_sigmoid(x)])
+ test_values = np.random.random((2, 5))
+ result = f([test_values])[0]
expected = hard_sigmoid(test_values)
self.assertAllClose(result, expected, rtol=1e-05)
def test_relu(self):
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.relu(x)])
- test_values = np.random.random((2, 5))
- result = f([test_values])[0]
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.relu(x)])
+ test_values = np.random.random((2, 5))
+ result = f([test_values])[0]
# No negative values in test values...
self.assertAllClose(result, test_values, rtol=1e-05)
- @test_util.run_deprecated_v1
def test_elu(self):
- with self.cached_session():
- x = keras.backend.placeholder(ndim=2)
- f = keras.backend.function([x], [keras.activations.elu(x, 0.5)])
- test_values = np.random.random((2, 5))
- result = f([test_values])[0]
- self.assertAllClose(result, test_values, rtol=1e-05)
- negative_values = np.array([[-1, -2]], dtype=keras.backend.floatx())
- result = f([negative_values])[0]
- true_result = (np.exp(negative_values) - 1) / 2
+ x = keras.backend.placeholder(ndim=2)
+ f = keras.backend.function([x], [keras.activations.elu(x, 0.5)])
+ test_values = np.random.random((2, 5))
+ result = f([test_values])[0]
+ self.assertAllClose(result, test_values, rtol=1e-05)
+ negative_values = np.array([[-1, -2]], dtype=keras.backend.floatx())
+ result = f([negative_values])[0]
+ true_result = (np.exp(negative_values) - 1) / 2
self.assertAllClose(result, true_result)
def test_tanh(self):
- with self.cached_session():
- test_values = np.random.random((2, 5))
- x = keras.backend.placeholder(ndim=2)
- exp = keras.activations.tanh(x)
- f = keras.backend.function([x], [exp])
- result = f([test_values])[0]
+ test_values = np.random.random((2, 5))
+ x = keras.backend.placeholder(ndim=2)
+ exp = keras.activations.tanh(x)
+ f = keras.backend.function([x], [exp])
+ result = f([test_values])[0]
expected = np.tanh(test_values)
self.assertAllClose(result, expected, rtol=1e-05)
def test_exponential(self):
- with self.cached_session():
- test_values = np.random.random((2, 5))
- x = keras.backend.placeholder(ndim=2)
- exp = keras.activations.exponential(x)
- f = keras.backend.function([x], [exp])
- result = f([test_values])[0]
+ test_values = np.random.random((2, 5))
+ x = keras.backend.placeholder(ndim=2)
+ exp = keras.activations.exponential(x)
+ f = keras.backend.function([x], [exp])
+ result = f([test_values])[0]
expected = np.exp(test_values)
self.assertAllClose(result, expected, rtol=1e-05)
diff --git a/tensorflow/python/keras/backend.py b/tensorflow/python/keras/backend.py
index 420c457..c93a716 100644
--- a/tensorflow/python/keras/backend.py
+++ b/tensorflow/python/keras/backend.py
@@ -568,6 +568,10 @@
Returns:
A list of available GPU devices.
"""
+ if ops.executing_eagerly_outside_functions():
+ # Returns names of devices directly.
+ return [name for name in context.list_devices() if 'GPU' in name]
+
global _LOCAL_DEVICES
if _LOCAL_DEVICES is None:
_LOCAL_DEVICES = get_session().list_devices()
@@ -2926,17 +2930,12 @@
def __init__(self, inputs, outputs, updates=None, name=None,
**session_kwargs):
updates = updates or []
- if not isinstance(inputs, (list, tuple)):
- raise TypeError('`inputs` to a Keras backend function '
- 'should be a list or tuple.')
- if not isinstance(outputs, (list, tuple)):
- raise TypeError('`outputs` of a Keras backend function '
- 'should be a list or tuple.')
if not isinstance(updates, (list, tuple)):
raise TypeError('`updates` in a Keras backend function '
'should be a list or tuple.')
- self.inputs = list(inputs)
- self.outputs = list(outputs)
+ self.inputs = nest.flatten(inputs)
+ self._outputs_structure = outputs
+ self.outputs = nest.flatten(outputs)
with ops.control_dependencies(self.outputs):
updates_ops = []
for update in updates:
@@ -3033,8 +3032,7 @@
self.fetch_callbacks[fetch](output)
def __call__(self, inputs):
- if not isinstance(inputs, (list, tuple)):
- raise TypeError('`inputs` should be a list or tuple.')
+ inputs = nest.flatten(inputs)
session = get_session()
feed_arrays = []
@@ -3077,7 +3075,8 @@
fetched = self._callable_fn(*array_vals,
run_metadata=self.run_metadata)
self._call_fetch_callbacks(fetched[-len(self._fetches):])
- return fetched[:len(self.outputs)]
+ return nest.pack_sequence_as(self._outputs_structure,
+ fetched[:len(self.outputs)])
class EagerExecutionFunction(object):
@@ -3093,17 +3092,12 @@
def __init__(self, inputs, outputs, updates=None, name=None):
updates = updates or []
- if not isinstance(inputs, (list, tuple)):
- raise TypeError('`inputs` to a Keras backend function '
- 'should be a list or tuple.')
- if not isinstance(outputs, (list, tuple)):
- raise TypeError('`outputs` of a Keras backend function '
- 'should be a list or tuple.')
if not isinstance(updates, (list, tuple)):
raise TypeError('`updates` in a Keras backend function '
'should be a list or tuple.')
- self.inputs = list(inputs)
- self.outputs = list(outputs)
+ self.inputs = nest.flatten(inputs)
+ self._outputs_structure = outputs
+ self.outputs = nest.flatten(outputs)
self.name = name
graph = get_graph()
@@ -3153,6 +3147,7 @@
x.op.inputs[0])
def __call__(self, inputs):
+ inputs = nest.flatten(inputs)
converted_inputs = []
for tensor, value in zip(self.inputs, inputs):
if value is None:
@@ -3169,7 +3164,8 @@
value = math_ops.cast(value, tensor.dtype)
converted_inputs.append(value)
outputs = self._graph_fn(*converted_inputs)
- return [x.numpy() for x in outputs]
+ return nest.pack_sequence_as(self._outputs_structure,
+ [x.numpy() for x in outputs])
@tf_export('keras.backend.function')
diff --git a/tensorflow/python/keras/backend_test.py b/tensorflow/python/keras/backend_test.py
index 6401e39..4b83f0b 100644
--- a/tensorflow/python/keras/backend_test.py
+++ b/tensorflow/python/keras/backend_test.py
@@ -1422,7 +1422,7 @@
decode_truth[i] == keras.backend.eval(decode_pred_tf[i])))
self.assertAllClose(log_prob_truth, log_prob_pred)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_ctc_batch_cost(self):
with self.cached_session():
label_lens = np.expand_dims(np.asarray([5, 4]), 1)
@@ -1695,6 +1695,39 @@
self.assertEqual(callback.times_called, 1)
self.assertEqual(callback.callback_result, 200)
+ @test_util.run_in_graph_and_eager_modes
+ def test_function_dict_outputs(self):
+ x_ph = keras.backend.placeholder(shape=(), name='x')
+ y_ph = keras.backend.placeholder(shape=(), name='y')
+ outputs = {'x*y': y_ph * x_ph, 'x*x': x_ph * x_ph}
+
+ f = keras.backend.function(inputs=[x_ph, y_ph], outputs=outputs)
+ x, y = 2., 5.
+ results = f([x, y])
+
+ self.assertEqual(results['x*y'], 10.)
+ self.assertEqual(results['x*x'], 4)
+
+ @test_util.run_in_graph_and_eager_modes
+ def test_function_dict_inputs(self):
+ placeholders = {
+ 'x': keras.backend.placeholder(shape=()),
+ 'y': keras.backend.placeholder(shape=())
+ }
+ outputs = [placeholders['x'] * placeholders['y']]
+
+ f = keras.backend.function(inputs=placeholders, outputs=outputs)
+ results = f({'x': 2., 'y': 3.})
+ self.assertEqual(results[0], 6.)
+
+ @test_util.run_in_graph_and_eager_modes
+ def test_function_single_input_output(self):
+ x_ph = keras.backend.placeholder(shape=(), name='x')
+ output = x_ph * x_ph
+ f = keras.backend.function(x_ph, output)
+ result = f(2.)
+ self.assertEqual(result, 4.)
+
def test_placeholder(self):
x = keras.backend.placeholder(shape=(3, 4))
self.assertEqual(x.get_shape().as_list(), [3, 4])
diff --git a/tensorflow/python/keras/callbacks_test.py b/tensorflow/python/keras/callbacks_test.py
index 6c9a382..4a65ade 100644
--- a/tensorflow/python/keras/callbacks_test.py
+++ b/tensorflow/python/keras/callbacks_test.py
@@ -403,7 +403,7 @@
float(keras.backend.get_value(
model.optimizer.lr)) - 0.01 / 4) < keras.backend.epsilon()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_ReduceLROnPlateau(self):
with self.cached_session():
np.random.seed(1337)
@@ -675,7 +675,7 @@
self.assertEqual(len(loss), 1)
self.assertEqual(loss[0], np.inf)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_TensorBoard(self):
np.random.seed(1337)
@@ -779,7 +779,7 @@
data_generator(True), len(x_train), epochs=2, callbacks=cbks)
assert os.path.exists(temp_dir)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_TensorBoard_multi_input_output(self):
np.random.seed(1337)
tmpdir = self.get_temp_dir()
@@ -851,7 +851,7 @@
callbacks=callbacks_factory(histogram_freq=1))
assert os.path.isdir(filepath)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_Tensorboard_histogram_summaries_in_test_function(self):
class FileWriterStub(object):
@@ -929,7 +929,7 @@
self.assertAllEqual(tsb.writer.steps_seen, [0, 1, 2, 3, 4, 5])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_Tensorboard_histogram_summaries_with_generator(self):
np.random.seed(1337)
tmpdir = self.get_temp_dir()
diff --git a/tensorflow/python/keras/constraints_test.py b/tensorflow/python/keras/constraints_test.py
index 4f674ea..92bc485 100644
--- a/tensorflow/python/keras/constraints_test.py
+++ b/tensorflow/python/keras/constraints_test.py
@@ -21,6 +21,7 @@
import numpy as np
from tensorflow.python import keras
+from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
@@ -35,6 +36,7 @@
return example_array
+@test_util.run_all_in_graph_and_eager_modes
class KerasConstraintsTest(test.TestCase):
def test_serialization(self):
@@ -49,54 +51,47 @@
assert fn.__class__ == ref_fn.__class__
def test_max_norm(self):
- with self.cached_session():
- array = get_example_array()
- for m in get_test_values():
- norm_instance = keras.constraints.max_norm(m)
- normed = norm_instance(keras.backend.variable(array))
- assert np.all(keras.backend.eval(normed) < m)
+ array = get_example_array()
+ for m in get_test_values():
+ norm_instance = keras.constraints.max_norm(m)
+ normed = norm_instance(keras.backend.variable(array))
+ assert np.all(keras.backend.eval(normed) < m)
- # a more explicit example
- norm_instance = keras.constraints.max_norm(2.0)
- x = np.array([[0, 0, 0], [1.0, 0, 0], [3, 0, 0], [3, 3, 3]]).T
- x_normed_target = np.array([[0, 0, 0], [1.0, 0, 0],
- [2.0, 0, 0],
- [2. / np.sqrt(3),
- 2. / np.sqrt(3),
- 2. / np.sqrt(3)]]).T
- x_normed_actual = keras.backend.eval(
- norm_instance(keras.backend.variable(x)))
- self.assertAllClose(x_normed_actual, x_normed_target, rtol=1e-05)
+ # a more explicit example
+ norm_instance = keras.constraints.max_norm(2.0)
+ x = np.array([[0, 0, 0], [1.0, 0, 0], [3, 0, 0], [3, 3, 3]]).T
+ x_normed_target = np.array(
+ [[0, 0, 0], [1.0, 0, 0], [2.0, 0, 0],
+ [2. / np.sqrt(3), 2. / np.sqrt(3), 2. / np.sqrt(3)]]).T
+ x_normed_actual = keras.backend.eval(
+ norm_instance(keras.backend.variable(x)))
+ self.assertAllClose(x_normed_actual, x_normed_target, rtol=1e-05)
def test_non_neg(self):
- with self.cached_session():
- non_neg_instance = keras.constraints.non_neg()
- normed = non_neg_instance(keras.backend.variable(get_example_array()))
- assert np.all(np.min(keras.backend.eval(normed), axis=1) == 0.)
+ non_neg_instance = keras.constraints.non_neg()
+ normed = non_neg_instance(keras.backend.variable(get_example_array()))
+ assert np.all(np.min(keras.backend.eval(normed), axis=1) == 0.)
def test_unit_norm(self):
- with self.cached_session():
- unit_norm_instance = keras.constraints.unit_norm()
- normalized = unit_norm_instance(
- keras.backend.variable(get_example_array()))
- norm_of_normalized = np.sqrt(
- np.sum(keras.backend.eval(normalized) ** 2, axis=0))
- # In the unit norm constraint, it should be equal to 1.
- difference = norm_of_normalized - 1.
- largest_difference = np.max(np.abs(difference))
- assert np.abs(largest_difference) < 10e-5
+ unit_norm_instance = keras.constraints.unit_norm()
+ normalized = unit_norm_instance(keras.backend.variable(get_example_array()))
+ norm_of_normalized = np.sqrt(
+ np.sum(keras.backend.eval(normalized)**2, axis=0))
+ # In the unit norm constraint, it should be equal to 1.
+ difference = norm_of_normalized - 1.
+ largest_difference = np.max(np.abs(difference))
+ assert np.abs(largest_difference) < 10e-5
def test_min_max_norm(self):
- with self.cached_session():
- array = get_example_array()
- for m in get_test_values():
- norm_instance = keras.constraints.min_max_norm(min_value=m,
- max_value=m * 2)
- normed = norm_instance(keras.backend.variable(array))
- value = keras.backend.eval(normed)
- l2 = np.sqrt(np.sum(np.square(value), axis=0))
- assert not l2[l2 < m]
- assert not l2[l2 > m * 2 + 1e-5]
+ array = get_example_array()
+ for m in get_test_values():
+ norm_instance = keras.constraints.min_max_norm(
+ min_value=m, max_value=m * 2)
+ normed = norm_instance(keras.backend.variable(array))
+ value = keras.backend.eval(normed)
+ l2 = np.sqrt(np.sum(np.square(value), axis=0))
+ assert not l2[l2 < m]
+ assert not l2[l2 > m * 2 + 1e-5]
if __name__ == '__main__':
diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index 8e35300..858fa76 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -1311,9 +1311,8 @@
def _loss_for_variable(v):
"""Creates a regularization loss `Tensor` for variable `v`."""
- with ops.colocate_with(v):
- with ops.name_scope(name + '/Regularizer'):
- regularization = regularizer(v)
+ with ops.name_scope(name + '/Regularizer'):
+ regularization = regularizer(v)
return regularization
if isinstance(variable, tf_variables.PartitionedVariable):
diff --git a/tensorflow/python/keras/engine/base_layer_test.py b/tensorflow/python/keras/engine/base_layer_test.py
index 798775b..fa4eb48 100644
--- a/tensorflow/python/keras/engine/base_layer_test.py
+++ b/tensorflow/python/keras/engine/base_layer_test.py
@@ -167,19 +167,26 @@
def test_mixing_keras_symbolic_tensors_and_eager_tensors(self):
x1 = keras.Input((3,))
x2 = array_ops.ones((3, 3))
- with self.assertRaisesRegexp(
- TypeError,
- 'mix computation of symbolic Tensors'):
- math_ops.matmul(x1, x2)
+ y = math_ops.matmul(x1, x2)
+ self.assertEqual(y.graph, keras.backend.get_graph())
+ fn = keras.backend.function(inputs=[x1], outputs=[y])
+ x_val = np.random.random((3, 3))
+ y_val = np.ones((3, 3))
+ self.assertAllClose(fn([x_val])[0],
+ np.matmul(x_val, y_val),
+ atol=1e-5)
def test_mixing_keras_symbolic_tensors_and_numpy_arrays(self):
- # For the time being we treat Numpy arrays as EagerTensors when mixing both.
x1 = keras.Input((3,))
x2 = np.ones((3, 3), dtype='float32')
- with self.assertRaisesRegexp(
- TypeError,
- 'mix computation of symbolic Tensors'):
- math_ops.matmul(x1, x2)
+ y = math_ops.matmul(x1, x2)
+ self.assertEqual(y.graph, keras.backend.get_graph())
+ fn = keras.backend.function(inputs=[x1], outputs=[y])
+ x_val = np.random.random((3, 3))
+ y_val = np.ones((3, 3))
+ self.assertAllClose(fn([x_val])[0],
+ np.matmul(x_val, y_val),
+ atol=1e-5)
if __name__ == '__main__':
diff --git a/tensorflow/python/keras/engine/distributed_training_utils.py b/tensorflow/python/keras/engine/distributed_training_utils.py
index d100182..32129af 100644
--- a/tensorflow/python/keras/engine/distributed_training_utils.py
+++ b/tensorflow/python/keras/engine/distributed_training_utils.py
@@ -199,11 +199,19 @@
# running ops.
if isinstance(callback, callbacks.TensorBoard):
if callback.__getattribute__('histogram_freq'):
- raise ValueError('histogram_freq in the TensorBoard callback is not '
- 'supported when using DistributionStrategy.')
+ logging.warning(
+ UserWarning(
+ '`histogram_freq` in the TensorBoard callback is not '
+ 'supported when using DistributionStrategy. Setting '
+ '`histogram_freq` to `0`.'))
+ callback.histogram_freq = 0
if callback.__getattribute__('write_grads'):
- raise ValueError('write_grads in the TensorBoard callback is not '
- 'supported when using DistributionStrategy.')
+ logging.warning(
+ UserWarning(
+ '`write_grads` in the TensorBoard callback is not supported '
+ 'when using DistributionStrategy. Setting `write_grads` '
+ 'to `False`.'))
+ callback.histogram_freq = False
def validate_distributed_dataset_inputs(distribution_strategy, x, y,
diff --git a/tensorflow/python/keras/engine/saving_test.py b/tensorflow/python/keras/engine/saving_test.py
index 6d9d9a2..f6ed3f4 100644
--- a/tensorflow/python/keras/engine/saving_test.py
+++ b/tensorflow/python/keras/engine/saving_test.py
@@ -30,8 +30,10 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
+from tensorflow.python.keras import optimizers
from tensorflow.python.keras.engine import saving
from tensorflow.python.keras.engine import training
+from tensorflow.python.keras.optimizer_v2 import rmsprop
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import random_ops
@@ -332,7 +334,6 @@
class TestWholeModelSaving(test.TestCase):
- @test_util.run_deprecated_v1
def test_sequential_model_saving(self):
if h5py is None:
self.skipTest('h5py required to run this test')
@@ -344,7 +345,7 @@
model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
model.compile(
loss=keras.losses.MSE,
- optimizer=keras.optimizers.RMSprop(lr=0.0001),
+ optimizer=rmsprop.RMSprop(lr=0.0001),
metrics=[
keras.metrics.categorical_accuracy,
keras.metrics.CategoricalAccuracy()
@@ -383,7 +384,10 @@
out = model.predict(x)
out2 = new_model.predict(x)
- self.assertAllClose(out, out2, atol=1e-05)
+
+ # TODO(b/120930751) This tolerance should be 1e-05,
+ # very concerning that its not.
+ self.assertAllClose(out, out2, atol=1e-03)
@test_util.run_deprecated_v1
def test_sequential_model_saving_without_input_shape(self):
@@ -635,8 +639,8 @@
os.close(fd)
os.remove(fname)
- @test_util.run_deprecated_v1
def test_saving_model_with_long_weights_names(self):
+ self.skipTest('b/120921503')
if h5py is None:
self.skipTest('h5py required to run this test')
@@ -756,14 +760,13 @@
class TestWeightSavingAndLoadingTFFormat(test.TestCase):
- @test_util.run_deprecated_v1
def test_keras_optimizer_warning(self):
graph = ops.Graph()
with graph.as_default(), self.session(graph):
model = keras.models.Sequential()
model.add(keras.layers.Dense(2, input_shape=(3,)))
model.add(keras.layers.Dense(3))
- model.compile(loss='mse', optimizer='adam', metrics=['acc'])
+ model.compile(loss='mse', optimizer=optimizers.Adam(), metrics=['acc'])
model._make_train_function()
temp_dir = self.get_temp_dir()
prefix = os.path.join(temp_dir, 'ckpt')
diff --git a/tensorflow/python/keras/engine/sequential_test.py b/tensorflow/python/keras/engine/sequential_test.py
index 001c3d7..10f69da 100644
--- a/tensorflow/python/keras/engine/sequential_test.py
+++ b/tensorflow/python/keras/engine/sequential_test.py
@@ -226,7 +226,7 @@
inner_model.trainable = True
self.assertEqual(len(model.trainable_weights), 4)
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_sequential_update_disabling(self):
val_a = np.random.random((10, 4))
val_out = np.random.random((10, 4))
diff --git a/tensorflow/python/keras/engine/topology_test.py b/tensorflow/python/keras/engine/topology_test.py
index 03bfd35..4071e2c 100644
--- a/tensorflow/python/keras/engine/topology_test.py
+++ b/tensorflow/python/keras/engine/topology_test.py
@@ -107,6 +107,7 @@
self.assertEqual(len(network.updates), 5)
self.assertEqual(len(network.get_updates_for(x4)), 2)
+ @test_util.run_v1_only('b/120545219')
def test_get_updates_bn(self):
x1 = input_layer_lib.Input(shape=(1,))
layer = keras.layers.BatchNormalization()
@@ -833,7 +834,7 @@
output_val_2 = m2.predict(x_val)
self.assertAllClose(output_val, output_val_2, atol=1e-6)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_explicit_training_argument(self):
with self.cached_session():
a = keras.layers.Input(shape=(2,))
diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index c902d92..38c8819 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -40,6 +40,7 @@
from tensorflow.python.keras.engine import training_generator
from tensorflow.python.keras.engine import training_utils
from tensorflow.python.keras.engine.network import Network
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow.python.keras.utils import data_utils
from tensorflow.python.keras.utils.generic_utils import slice_arrays
from tensorflow.python.keras.utils.losses_utils import squeeze_or_expand_dimensions
@@ -128,307 +129,6 @@
self.run_eagerly = None
- def _set_sample_weight_attributes(self, sample_weight_mode,
- skip_target_weighing_indices):
- """Sets sample weight related attributes on the model."""
- sample_weights, sample_weight_modes = training_utils.prepare_sample_weights(
- self.output_names, sample_weight_mode, skip_target_weighing_indices)
- self.sample_weights = sample_weights
- self.sample_weight_modes = sample_weight_modes
- self._feed_sample_weight_modes = [
- sample_weight_modes[i]
- for i in range(len(self.outputs))
- if i not in skip_target_weighing_indices
- ]
- self._feed_sample_weights = [
- sample_weights[i]
- for i in range(len(sample_weights))
- if i not in skip_target_weighing_indices
- ]
-
- def _cache_output_metric_attributes(self, metrics, weighted_metrics):
- """Caches metric name and function attributes for every model output."""
- output_shapes = [
- None if output is None else output.get_shape().as_list()
- for output in self.outputs
- ]
- self._per_output_metrics = training_utils.collect_per_output_metric_info(
- metrics, self.output_names, output_shapes, self.loss_functions)
- self._per_output_weighted_metrics = \
- training_utils.collect_per_output_metric_info(
- weighted_metrics, self.output_names, output_shapes,
- self.loss_functions, self.sample_weights)
-
- def _add_unique_metric_name(self, metric_name, output_index):
- """Makes the metric name unique and adds it to the model's metric name list.
-
- If there are multiple outputs for which the metrics are calculated, the
- metric names have to be made unique by appending an integer.
-
- Arguments:
- metric_name: Metric name that corresponds to the metric specified by the
- user. For example: 'acc'.
- output_index: The index of the model output for which the metric name is
- being added.
-
- Returns:
- string, name of the model's unique metric name
- """
- if len(self.output_names) > 1:
- metric_name = '%s_%s' % (self.output_names[output_index], metric_name)
- j = 1
- base_metric_name = metric_name
- while metric_name in self._compile_metrics_names:
- metric_name = '%s_%d' % (base_metric_name, j)
- j += 1
-
- return metric_name
-
- @property
- def metrics(self):
- """Returns the model's metrics added using `compile`, `add_metric` APIs."""
- metrics = []
- if self._is_compiled:
- metrics += self._compile_stateful_metric_functions
- return metrics + super(Model, self).metrics
-
- @property
- def metrics_names(self):
- """Returns the model's display labels for all outputs."""
- metrics_names = []
- if self._is_compiled:
- metrics_names += self._compile_metrics_names # Includes names of losses.
-
- # Add metric names from layers.
- for layer in self.layers:
- metrics_names += [m.name for m in layer._metrics] # pylint: disable=protected-access
- metrics_names += [m.name for m in self._metrics]
- return metrics_names
-
- @property
- def _all_metrics_tensors(self):
- """Returns the network's symbolic metric tensors."""
- metrics_tensors = {}
- if self._is_compiled:
- metrics_tensors.update(self._compile_metrics_tensors)
- metrics_tensors.update(super(Model, self)._all_metrics_tensors)
- return metrics_tensors
-
- @property
- def _all_stateful_metrics_tensors(self):
- """Returns the network's symbolic metric tensors."""
- metrics_tensors = {}
- if self._is_compiled:
- metrics_tensors.update(self._compile_stateful_metrics_tensors)
- metrics_tensors.update(super(Model, self)._all_metrics_tensors)
- return metrics_tensors
-
- def _init_metric_attributes(self):
- """Initialized model metric attributes."""
- # List of all metric names in the model.
- self._compile_metrics_names = ['loss']
- # List of stateful metric functions. Used for resetting metric state during
- # training/eval.
- # This includes loss functions when there are multiple outputs.
- self._compile_stateful_metric_functions = []
- # Dict of all aggregated metric result tensors. This includes aggregated
- # loss result tensors when there are multiple outputs.
- self._compile_stateful_metrics_tensors = {}
- # Dict of all metric result tensors (aggregated or not - based on the
- # values given in compile.). This includes aggregated loss result tensors
- # when there are multiple outputs.
- self._compile_metrics_tensors = {}
-
- def _set_per_output_metric_attributes(self, metrics_dict, output_index):
- """Sets the metric attributes on the model for the given output.
-
- Arguments:
- metrics_dict: A dict with metric names as keys and metric fns as values.
- output_index: The index of the model output for which the metric
- attributes are added.
-
- Returns:
- Metrics dict updated with unique metric names as keys.
- """
- updated_metrics_dict = collections.OrderedDict()
- for metric_name, (metric_fn, stateful_metric_fn) in metrics_dict.items():
- metric_name = self._add_unique_metric_name(metric_name, output_index)
- updated_metrics_dict[metric_name] = (metric_fn, stateful_metric_fn)
- # Keep track of metric name, function and stateful function.
- self._compile_metrics_names.append(metric_name)
- self._compile_stateful_metric_functions.append(stateful_metric_fn)
- return updated_metrics_dict
-
- def _set_metric_attributes(self, outputs, skip_target_indices=None):
- """Sets the metric attributes on the model for all the model outputs."""
- skip_target_indices = skip_target_indices or []
- updated_per_output_metrics = []
- updated_per_output_weighted_metrics = []
- for i in range(len(outputs)):
- if i in skip_target_indices:
- updated_per_output_metrics.append(self._per_output_metrics[i])
- updated_per_output_weighted_metrics.append(
- self._per_output_weighted_metrics[i])
- continue
- updated_per_output_metrics.append(
- self._set_per_output_metric_attributes(self._per_output_metrics[i],
- i))
- updated_per_output_weighted_metrics.append(
- self._set_per_output_metric_attributes(
- self._per_output_weighted_metrics[i], i))
-
- self._per_output_metrics = updated_per_output_metrics
- self._per_output_weighted_metrics = updated_per_output_weighted_metrics
-
- def _handle_per_output_metrics(self,
- metrics_dict,
- y_true,
- y_pred,
- mask,
- weights=None,
- return_stateful_result=True):
- """Calls metric functions for a single output.
-
- Arguments:
- metrics_dict: A dict with metric names as keys and metric fns as values.
- y_true: Target output.
- y_pred: Predicted output.
- mask: Computed mask value for the current output.
- weights: Weights to be applied on the current output.
- return_stateful_result: Boolean, indicates whether the stateful
- (aggregated)/stateless metric result should be returned.
-
- Returns:
- A list of metric result tensors.
- """
- metric_results = []
- for metric_name, (metric_fn, stateful_fn) in metrics_dict.items():
- with K.name_scope(metric_name):
-
- def _call_stateful_fn(fn):
- return training_utils.call_metric_function(
- fn, y_true, y_pred, weights=weights, mask=mask)
-
- def _call_stateless_fn(fn):
- weighted_metric_fn = training_utils.weighted_masked_objective(fn)
- return weighted_metric_fn(y_true, y_pred, weights=weights, mask=mask)
-
- def _track_metric_tensors(name, stateless_result, stateful_result):
- self._compile_metrics_tensors[name] = stateless_result
- self._compile_stateful_metrics_tensors[name] = stateful_result
-
- if isinstance(metric_fn, metrics_module.Metric):
- # If the given metric fn is stateful, call the fn and return result.
- metric_result = _call_stateful_fn(metric_fn)
- metric_results.append(metric_result)
- if not self.run_eagerly:
- _track_metric_tensors(metric_name, metric_result, metric_result)
- elif self.run_eagerly:
- # In eager mode, if the given metric fn is not stateful, we invoke the
- # given fn or its stateful version based on the given flag.
- if return_stateful_result:
- metric_result = _call_stateful_fn(stateful_fn)
- else:
- metric_result = _call_stateless_fn(metric_fn)
- metric_results.append(metric_result)
- else:
- # In graph mode, we build the sub-graph for both the stateful and the
- # stateless fns.
- stateful_metric_result = _call_stateful_fn(stateful_fn)
- metric_result = _call_stateless_fn(metric_fn)
- _track_metric_tensors(metric_name, metric_result,
- stateful_metric_result)
-
- return metric_results
-
- def _handle_metrics(self,
- outputs,
- skip_target_indices=None,
- targets=None,
- sample_weights=None,
- masks=None,
- return_stateful_result=True):
- """Handles calling metric functions.
-
- Arguments:
- outputs: List of outputs (predictions).
- skip_target_indices: Optional. List of target ids to skip.
- targets: List of targets.
- sample_weights: Optional list of sample weight arrays.
- masks: List of computed output mask values.
- return_stateful_result: Boolean, indicates whether the stateful
- (aggregated)/stateless metric result should be returned.
-
- Returns:
- A list of metric result tensors.
- """
- skip_target_indices = skip_target_indices or []
- metric_results = []
- with K.name_scope('metrics'):
- # Invoke all metrics added using `compile`.
- for i in range(len(outputs)):
- if i in skip_target_indices:
- continue
- output = outputs[i] if outputs else None
- target = targets[i] if targets else None
- output_mask = masks[i] if masks else None
- metric_results.extend(
- self._handle_per_output_metrics(
- self._per_output_metrics[i],
- target,
- output,
- output_mask,
- return_stateful_result=return_stateful_result))
- metric_results.extend(
- self._handle_per_output_metrics(
- self._per_output_weighted_metrics[i],
- target,
- output,
- output_mask,
- weights=sample_weights[i],
- return_stateful_result=return_stateful_result))
-
- # Add metric results from the `add_metric` metrics in eager mode.
- if context.executing_eagerly():
- for m in self.metrics:
- if m not in self._compile_stateful_metric_functions:
- metric_results.append(m.result())
- return metric_results
-
- @property
- def run_eagerly(self):
- """Settable attribute indicating whether the model should run eagerly.
-
- Running eagerly means that your model will be run step by step,
- like Python code. Your model might run slower, but it should become easier
- for you to debug it by stepping into individual layer calls.
-
- By default, we will attempt to compile your model to a static graph to
- deliver the best execution performance.
-
- Returns:
- Boolean, whether the model should run eagerly.
- """
- if self._run_eagerly is True and not context.executing_eagerly():
- raise ValueError('You can only set `run_eagerly=True` if eager execution '
- 'is enabled.')
- if self._static_graph_friendly:
- if self._run_eagerly is None:
- return False
- else:
- return self._run_eagerly
- else:
- if self._run_eagerly is False:
- # TODO(fchollet): consider using py_func to enable this.
- raise ValueError('Your model contains layers that can only be '
- 'successfully run in eager execution. '
- 'You cannot set `run_eagerly=False`.')
- return context.executing_eagerly()
-
- @run_eagerly.setter
- def run_eagerly(self, value):
- self._run_eagerly = value
-
@checkpointable.no_automatic_dependency_tracking
def compile(self,
optimizer,
@@ -444,12 +144,11 @@
Arguments:
optimizer: String (name of optimizer) or optimizer instance.
- See [optimizers](/api_docs/python/tf/keras/optimizers).
+ See `tf.keras.optimizers`.
loss: String (name of objective function) or objective function.
- See [losses](/api_docs/python/tf/losses).
- If the model has multiple outputs, you can use a different loss
- on each output by passing a dictionary or a list of losses.
- The loss value that will be minimized by the model
+ See `tf.losses`. If the model has multiple outputs, you can use a
+ different loss on each output by passing a dictionary or a list of
+ losses. The loss value that will be minimized by the model
will then be the sum of all individual losses.
metrics: List of metrics to be evaluated by the model
during training and testing.
@@ -492,12 +191,14 @@
"""
run_eagerly = kwargs.pop('run_eagerly', None)
self._run_eagerly = run_eagerly
+ optimizer = optimizers.get(optimizer)
# Validate that arguments passed by the user to `compile` are supported by
# DistributionStrategy.
if distribute:
- if not isinstance(
- optimizer, (tf_optimizer_module.Optimizer, optimizers.TFOptimizer)):
+ if not isinstance(optimizer,
+ (tf_optimizer_module.Optimizer, optimizers.TFOptimizer,
+ optimizer_v2.OptimizerV2)):
raise NotImplementedError(
'optimizer must be an instance of '
'tf.train.Optimizer, not a %s' % type(optimizer))
@@ -513,13 +214,14 @@
loss = loss or {}
if self.run_eagerly and not isinstance(
- optimizer, (tf_optimizer_module.Optimizer, optimizers.TFOptimizer)):
+ optimizer, (tf_optimizer_module.Optimizer, optimizers.TFOptimizer,
+ optimizer_v2.OptimizerV2)):
raise ValueError(
'When running a model in eager execution, the optimizer must be an '
'instance of tf.train.Optimizer. Received: '
'%s' % optimizer)
- self.optimizer = optimizers.get(optimizer)
+ self.optimizer = optimizer
# We've disabled automatic dependency tracking for this method, but do want
# to add a checkpoint dependency on the optimizer if it's checkpointable.
if isinstance(self.optimizer, checkpointable.CheckpointableBase):
@@ -814,740 +516,60 @@
trainable_weights = self.trainable_weights
self._collected_trainable_weights = trainable_weights
- def _check_trainable_weights_consistency(self):
- """Check trainable weights count consistency.
+ @property
+ def metrics(self):
+ """Returns the model's metrics added using `compile`, `add_metric` APIs."""
+ metrics = []
+ if self._is_compiled:
+ metrics += self._compile_stateful_metric_functions
+ return metrics + super(Model, self).metrics
- This will raise a warning if `trainable_weights` and
- `_collected_trainable_weights` are inconsistent (i.e. have different
- number of parameters).
- Inconsistency will typically arise when one modifies `model.trainable`
- without calling `model.compile` again.
- """
- if not hasattr(self, '_collected_trainable_weights'):
- return
+ @property
+ def metrics_names(self):
+ """Returns the model's display labels for all outputs."""
+ metrics_names = []
+ if self._is_compiled:
+ metrics_names += self._compile_metrics_names # Includes names of losses.
- if len(self.trainable_weights) != len(self._collected_trainable_weights):
- logging.log_first_n(
- logging.WARN, 'Discrepancy between trainable weights and collected'
- ' trainable weights, did you set `model.trainable`'
- ' without calling `model.compile` after ?', 1)
+ # Add metric names from layers.
+ for layer in self.layers:
+ metrics_names += [m.name for m in layer._metrics] # pylint: disable=protected-access
+ metrics_names += [m.name for m in self._metrics]
+ return metrics_names
- def _make_train_function_helper(self, fn_name, outputs, metric_updates=None):
- if not hasattr(self, fn_name):
- raise RuntimeError('You must compile your model before using it.')
- self._check_trainable_weights_consistency()
- if getattr(self, fn_name) is None:
- inputs = (self._feed_inputs +
- self._feed_targets +
- self._feed_sample_weights)
- if not isinstance(K.symbolic_learning_phase(), int):
- inputs += [K.symbolic_learning_phase()]
+ @property
+ def run_eagerly(self):
+ """Settable attribute indicating whether the model should run eagerly.
- with K.get_graph().as_default():
- with K.name_scope('training'):
- with K.name_scope(self.optimizer.__class__.__name__):
- # Training updates
- updates = self.optimizer.get_updates(
- params=self._collected_trainable_weights, loss=self.total_loss)
- # Unconditional updates
- updates += self.get_updates_for(None)
- # Conditional updates relevant to this model
- updates += self.get_updates_for(self.inputs)
- # Add stateful metrics updates.
- if metric_updates is not None:
- updates += metric_updates
+ Running eagerly means that your model will be run step by step,
+ like Python code. Your model might run slower, but it should become easier
+ for you to debug it by stepping into individual layer calls.
- with K.name_scope('training'):
- # Gets loss and metrics. Updates weights at each call.
- fn = K.function(
- inputs,
- outputs,
- updates=updates,
- name='train_function',
- **self._function_kwargs)
- setattr(self, fn_name, fn)
-
- def _make_train_function(self):
- metrics_tensors = [
- self._all_metrics_tensors[m] for m in self.metrics_names[1:]
- ]
- self._make_train_function_helper('train_function',
- [self.total_loss] + metrics_tensors)
-
- def _make_fit_function(self):
- # TODO(psv/anjalisridhar): Remove updates after we fix b/118841692
- # Stateful metrics updates
- metric_updates = []
- for m in self.metrics:
- metric_updates += m.updates
-
- metrics_tensors = [
- self._all_stateful_metrics_tensors[m] for m in self.metrics_names[1:]
- ]
- self._make_train_function_helper(
- '_fit_function', [self.total_loss] + metrics_tensors, metric_updates)
-
- def _make_test_function_helper(self, fn_name, outputs, metric_updates=None):
- if not hasattr(self, fn_name):
- raise RuntimeError('You must compile your model before using it.')
- if getattr(self, fn_name) is None:
- inputs = (self._feed_inputs +
- self._feed_targets +
- self._feed_sample_weights)
-
- with K.name_scope('evaluation'):
- updates = self.state_updates
- # Add stateful metrics updates.
- if metric_updates is not None:
- updates += metric_updates
- # Return loss and metrics, no gradient updates.
- # Does update the network states.
- fn = K.function(
- inputs,
- outputs,
- updates=updates,
- name='test_function',
- **self._function_kwargs)
- setattr(self, fn_name, fn)
-
- def _make_test_function(self):
- metrics_tensors = [
- self._all_metrics_tensors[m] for m in self.metrics_names[1:]
- ]
- self._make_test_function_helper('test_function',
- [self.total_loss] + metrics_tensors)
-
- def _make_eval_function(self):
- metrics_tensors = [
- self._all_stateful_metrics_tensors[m] for m in self.metrics_names[1:]
- ]
- self._make_test_function_helper('_eval_function',
- [self.total_loss] + metrics_tensors)
-
- def _make_predict_function(self):
- if not hasattr(self, 'predict_function'):
- self.predict_function = None
- if self.predict_function is None:
- inputs = self._feed_inputs
- # Gets network outputs. Does not update weights.
- # Does update the network states.
- kwargs = getattr(self, '_function_kwargs', {})
- with K.name_scope('predict'):
- self.predict_function = K.function(
- inputs,
- self.outputs,
- updates=self.state_updates,
- name='predict_function',
- **kwargs)
-
- def _make_execution_function(self, mode):
- if mode == 'train':
- self._make_fit_function()
- return self._fit_function
- if mode == 'test':
- self._make_eval_function()
- return self._eval_function
- if mode == 'predict':
- self._make_predict_function()
- return self.predict_function
-
- def _get_iterator_get_next_tensors(self, iterator):
- get_next_op = self._iterator_get_next.get(iterator, None)
- if get_next_op is None:
- get_next_op = iterator.get_next()
- self._iterator_get_next[iterator] = get_next_op
- return get_next_op
-
- def _distribution_standardize_user_data(self,
- x,
- y=None,
- sample_weight=None,
- class_weight=None,
- batch_size=None,
- check_steps=False,
- steps_name='steps',
- steps=None,
- validation_split=0,
- shuffle=False):
- """Runs validation checks on input and target data passed by the user.
-
- This is called when using DistributionStrategy to train, evaluate or serve
- the model.
-
- Args:
- x: Input data. A numpy array or `tf.data` dataset.
- y: Target data. A numpy array or None if x is a `tf.data` dataset.
- sample_weight: An optional sample-weight array passed by the user to
- weight the importance of each sample in `x`.
- class_weight: An optional class-weight array by the user to
- weight the importance of samples in `x` based on the class they belong
- to, as conveyed by `y`.
- batch_size: Integer batch size. If provided, it is used to run additional
- validation checks on stateful models.
- check_steps: boolean, True if we want to check for validity of `steps` and
- False, otherwise.
- steps_name: The public API's parameter name for `steps`.
- steps: Integer or `None`. Total number of steps (batches of samples) to
- execute.
- validation_split: Float between 0 and 1.
- Fraction of the training data to be used as validation data.
- shuffle: Boolean whether to shuffle the training data before each epoch.
+ By default, we will attempt to compile your model to a static graph to
+ deliver the best execution performance.
Returns:
- Iterator for reading the dataset `x`.
-
- Raises:
- ValueError: In case of invalid user-provided data.
- RuntimeError: If the model was never compiled.
+ Boolean, whether the model should run eagerly.
"""
- if class_weight:
- raise NotImplementedError('`class_weight` is currently not supported '
- 'when using DistributionStrategy.')
-
- if (sample_weight is not None and sample_weight.all() and
- distributed_training_utils.is_tpu_strategy(
- self._distribution_strategy)):
- raise NotImplementedError('`sample_weight` is currently not supported '
- 'when using TPUStrategy.')
-
- # Validates `steps` argument right at the beginning since we use it to
- # construct the dataset object.
- # TODO(anjalisridhar): Remove this check once we refactor the
- # _standardize_user_data code path. This check is already present elsewhere
- # in the codebase.
- if check_steps and isinstance(x, dataset_ops.DatasetV2) and steps is None:
- raise ValueError('When using Datasets as input, '
- 'you should specify the `{steps_name}` argument.'
- .format(steps_name=steps_name))
-
- first_x_value = nest.flatten(x)[0]
- if isinstance(first_x_value, np.ndarray):
- # We need to use the drop_remainder argument to allow for a static
- # input shape which is required for TPUs.
- drop_remainder = self._distribution_strategy.require_static_shapes
- if y is not None:
- var_x = distributed_training_utils.get_var_for_numpy(
- self._distribution_strategy, x)
- var_y = distributed_training_utils.get_var_for_numpy(
- self._distribution_strategy, y)
- if sample_weight is not None:
- var_sample_weights = distributed_training_utils.get_var_for_numpy(
- self._distribution_strategy, sample_weight)
-
- x = dataset_ops.Dataset.from_tensor_slices((var_x, var_y,
- var_sample_weights))
- else:
- x = dataset_ops.Dataset.from_tensor_slices((var_x, var_y))
-
- x = dataset_ops.Dataset.from_tensor_slices((var_x, var_y))
- if shuffle:
- # 1024 is a good buffer size since it is much larger than the average
- # batch size provided by the user and provides sufficient randomness.
- # One thing to keep in mind is the memory usage based on the size of
- # each sample.
- x = x.shuffle(1024)
- x = x.repeat()
- x = x.batch(batch_size, drop_remainder=drop_remainder)
- y = None
- sample_weight = None
+ if self._run_eagerly is True and not context.executing_eagerly():
+ raise ValueError('You can only set `run_eagerly=True` if eager execution '
+ 'is enabled.')
+ if self._static_graph_friendly:
+ if self._run_eagerly is None:
+ return False
else:
- # This case is for the predict call where the dataset only contains
- # inputs and no targets, i.e. it does not return a tuple
- var_x = distributed_training_utils.get_var_for_numpy(
- self._distribution_strategy, x)
- x = dataset_ops.Dataset.from_tensor_slices(var_x)
- x = x.batch(batch_size, drop_remainder=drop_remainder)
-
- assert isinstance(x, dataset_ops.DatasetV2)
-
- with self._distribution_strategy.scope():
- iterator = self._distribution_strategy.make_dataset_iterator(x)
- init_op = iterator.initialize()
- if not context.executing_eagerly():
- K.get_session().run(init_op)
-
- training_utils.validate_iterator_input(x, y, sample_weight,
- validation_split)
- return iterator
-
- def _standardize_user_data(self,
- x,
- y=None,
- sample_weight=None,
- class_weight=None,
- batch_size=None,
- check_steps=False,
- steps_name='steps',
- steps=None,
- validation_split=0,
- shuffle=False):
- """Runs validation checks on input and target data passed by the user.
-
- Also standardizes the data to lists of arrays, in order.
-
- Also builds and compiles the model on the fly if it is a subclassed model
- that has never been called before (and thus has no inputs/outputs).
-
- This is a purely internal method, subject to refactoring at any time.
-
- Args:
- x: Input data. It could be:
- - A Numpy array (or array-like), or a list of arrays
- (in case the model has multiple inputs).
- - A TensorFlow tensor, or a list of tensors
- (in case the model has multiple inputs).
- - A dict mapping input names to the corresponding array/tensors,
- if the model has named inputs.
- - A `tf.data` dataset or a dataset iterator.
- y: Target data. Like the input data `x`,
- it could be either Numpy array(s) or TensorFlow tensor(s).
- It should be consistent with `x` (you cannot have Numpy inputs and
- tensor targets, or inversely). If `x` is a dataset or a
- dataset iterator, `y` should not be specified
- (since targets will be obtained from the iterator).
- sample_weight: An optional sample-weight array passed by the user to
- weight the importance of each sample in `x`.
- class_weight: An optional class-weight array by the user to
- weight the importance of samples in `x` based on the class they belong
- to, as conveyed by `y`.
- batch_size: Integer batch size. If provided, it is used to run additional
- validation checks on stateful models.
- check_steps: boolean, True if we want to check for validity of `steps` and
- False, otherwise. For example, when we are standardizing one batch of
- data for train_on_batch/predict_on_batch/test_on_batch APIs, `steps`
- value is not required and we should not check for its validity in these
- cases.
- steps_name: The public API's parameter name for `steps`.
- steps: Integer or `None`. Total number of steps (batches of samples) to
- execute.
- validation_split: Float between 0 and 1.
- Fraction of the training data to be used as validation data.
- shuffle: Boolean whether to shuffle the training data before each epoch.
-
- Returns:
- A tuple of 3: inputs (arrays or dicts, depending on whether `x` was a dict
- or not), target arrays, sample-weight arrays.
- If the model's input and targets are symbolic, these lists are empty
- (since the model takes no user-provided data, instead the data comes
- from the symbolic inputs/targets).
-
- Raises:
- ValueError: In case of invalid user-provided data.
- RuntimeError: If the model was never compiled.
- """
- if self._distribution_strategy:
- iterator = self._distribution_standardize_user_data(
- x,
- y,
- sample_weight=sample_weight,
- class_weight=class_weight,
- batch_size=batch_size,
- check_steps=check_steps,
- steps_name=steps_name,
- steps=steps,
- validation_split=validation_split,
- shuffle=shuffle)
- return iterator, None, None
-
- if isinstance(x, dataset_ops.DatasetV2):
- if context.executing_eagerly():
- x = iter(x)
- else:
- if x in self._dataset_iterator_cache:
- x = self._dataset_iterator_cache[x]
- else:
- iterator = dataset_ops.make_initializable_iterator(x)
- self._dataset_iterator_cache[x] = iterator
- x = iterator
- K.get_session().run(x.initializer)
-
- # Validates `steps` argument based on x's type.
- if check_steps:
- training_utils.check_steps_argument(x, steps, steps_name)
-
- is_x_eager_iterator = isinstance(x, iterator_ops.EagerIterator)
- is_x_iterator = isinstance(x, iterator_ops.Iterator)
-
- # Validate user inputs when data is given as a dataset or dataset iterator.
- if is_x_iterator or is_x_eager_iterator:
- training_utils.validate_iterator_input(x, y, sample_weight,
- validation_split)
-
- # For eager iterators, when we have to process multiple batches of samples,
- # we will standardize the data when we actually loop over iterator and get
- # the batches. For now, we just return the iterator as is.
- if is_x_eager_iterator:
- return x, y, sample_weight
-
- # If input data is a dataset iterator in graph mode or if it is an eager
- # iterator and only one batch of samples is required, we fetch the data
- # tensors from the iterator and then standardize them.
- if is_x_iterator or is_x_eager_iterator:
- try:
- if is_x_iterator:
- next_element = self._get_iterator_get_next_tensors(x)
- else:
- next_element = x.get_next()
- except errors.OutOfRangeError:
- raise RuntimeError('Your dataset iterator ran out of data; '
- 'Make sure that your dataset can generate '
- 'required number of samples.')
-
- if isinstance(next_element, (list, tuple)):
- if len(next_element) not in [2, 3]:
- raise ValueError(
- 'Please provide model inputs as a list or tuple of 2 or 3'
- 'elements: (input, target) or (input, target, sample_weights)'
- 'Received %s' % next_element)
- if len(next_element) == 2:
- x, y = next_element
- else:
- x, y, sample_weight = next_element
- else:
- x = next_element
- x, y, sample_weights = self._standardize_weights(
- x, y, sample_weight, class_weight, batch_size, is_x_iterator)
- return x, y, sample_weights
-
- def _standardize_weights(self,
- x,
- y,
- sample_weight=None,
- class_weight=None,
- batch_size=None,
- from_iterator=False):
- """Standardize input data, target data, and weight values.
-
- This method reformats all data passed to the model to an ordered list of
- array/tensors, matching the order expected by the model. This also validates
- the input and target data shapes.
-
- Args:
- x: Input data. It could be:
- - A Numpy array (or array-like), or a list of arrays
- (in case the model has multiple inputs).
- - A TensorFlow tensor, or a list of tensors
- (in case the model has multiple inputs).
- - A dict mapping input names to the corresponding array/tensors,
- if the model has named inputs.
- x cannot not be an iterator.
- y: Target data. Like the input data `x`,
- it could be either Numpy array(s) or TensorFlow tensor(s).
- It should be consistent with `x` (you cannot have Numpy inputs and
- tensor targets, or inversely).
- sample_weight: An optional sample-weight array passed by the user to
- weight the importance of each sample in `x`.
- class_weight: An optional class-weight array by the user to
- weight the importance of samples in `x` based on the class they belong
- to, as conveyed by `y`.
- batch_size: Integer batch size. If provided, it is used to run additional
- validation checks on stateful models.
- from_iterator: Whether x and y were obtained from an iterator.
-
- Returns:
- Tuple of standardized data that will be fed to the model:
- (input data, target data, sample weights)
-
- Raises:
- RuntimeError: If target data is provided, but the model has not yet been
- compiled.
- ValueError: If the input data, target data, and batch size have invalid
- shapes or formats (e.g. the model expects input to be a list of three
- tensors, but x is a list with two tensors). Error is also raised if the
- input and target data are not both arrays or tensors.
- """
- # TODO(sourabhbajaj): Split input validation from weight standardization.
- if sample_weight is not None and class_weight is not None:
- logging.warning(
- 'Received both a `sample_weight` and `class_weight` argument. '
- 'The `class_weight` argument will be ignored.')
- # First, we build/compile the model on the fly if necessary.
- all_inputs = []
- is_build_called = False
- is_compile_called = False
- # Whether this is a subclassed model that expects dictionary inputs
- # rather than list inputs (e.g. FeatureColumn-based models).
- dict_inputs = False
- if not self.inputs:
- # We need to use `x` to set the model inputs.
- # We type-check that `x` and `y` are either single arrays
- # or lists of arrays.
- if isinstance(x, (list, tuple)):
- if not all(isinstance(v, np.ndarray) or
- tensor_util.is_tensor(v) for v in x):
- raise ValueError('Please provide as model inputs either a single '
- 'array or a list of arrays. You passed: x=' + str(x))
- all_inputs += list(x)
- elif isinstance(x, dict):
- dict_inputs = True
- keys = sorted(x.keys())
- all_inputs = [x[k] for k in keys]
- else:
- if not isinstance(x, np.ndarray) and not tensor_util.is_tensor(x):
- raise ValueError('Please provide as model inputs either a single '
- 'array or a list of arrays. You passed: x=' + str(x))
- all_inputs.append(x)
-
- # Build the model using the retrieved inputs (value or symbolic).
- # If values or generated from a dataset, then in symbolic-mode
- # placeholders will be created to match the value shapes.
- if not self.inputs:
- is_build_called = True
- if from_iterator:
- cast_inputs = nest.map_structure(lambda v: v.shape, x)
- elif training_utils.has_tensors(x):
- cast_inputs = training_utils.cast_if_floating_dtype(x)
- else:
- cast_inputs = x
- self._set_inputs(cast_inputs)
+ return self._run_eagerly
else:
- dict_inputs = isinstance(self.inputs, dict)
- if dict_inputs and context.executing_eagerly():
- # No support for graph functions when the model expects dictionary inputs
- # (i.e. FeatureColumn-based models).
- self.run_eagerly = True
+ if self._run_eagerly is False:
+ # TODO(fchollet): consider using py_func to enable this.
+ raise ValueError('Your model contains layers that can only be '
+ 'successfully run in eager execution. '
+ 'You cannot set `run_eagerly=False`.')
+ return context.executing_eagerly()
- if y is not None:
- if not self.optimizer:
- raise RuntimeError('You must compile a model before '
- 'training/testing. '
- 'Use `model.compile(optimizer, loss)`.')
- if not self._is_compiled:
- # On-the-fly compilation of the model.
- # We need to use `y` to set the model targets.
- if training_utils.has_tensors(y):
- y = training_utils.cast_if_floating_dtype(y)
- if isinstance(y, (list, tuple)):
- if not all(isinstance(v, np.ndarray) or
- tensor_util.is_tensor(v) for v in y):
- raise ValueError('Please provide as model targets either a single '
- 'array or a list of arrays. '
- 'You passed: y=' + str(y))
- all_inputs += list(y)
- elif isinstance(y, dict):
- raise ValueError('Please do not pass a dictionary as model targets.')
- else:
- if not isinstance(y, np.ndarray) and not tensor_util.is_tensor(y):
- raise ValueError('Please provide as model targets either a single '
- 'array or a list of arrays. '
- 'You passed: y=' + str(y))
- all_inputs.append(y)
-
- # Typecheck that all inputs are *either* value *or* symbolic.
- # TODO(fchollet): this check could be removed in Eager mode?
- if any(tensor_util.is_tensor(v) for v in all_inputs):
- if not all(tensor_util.is_tensor(v) for v in all_inputs):
- raise ValueError('Do not pass inputs that mix Numpy arrays and '
- 'TensorFlow tensors. '
- 'You passed: x=' + str(x) + '; y=' + str(y))
-
- if self.run_eagerly or from_iterator:
- target_tensors = None
- else:
- # Handle target tensors if any passed.
- if not isinstance(y, (list, tuple)):
- y = [y]
- target_tensors = [v for v in y if _is_symbolic_tensor(v)]
- is_compile_called = True
- self.compile(
- optimizer=self.optimizer,
- loss=self.loss,
- metrics=self._compile_metrics,
- weighted_metrics=self._compile_weighted_metrics,
- loss_weights=self.loss_weights,
- target_tensors=target_tensors,
- run_eagerly=self.run_eagerly)
-
- # In graph mode, if we had just set inputs and targets as symbolic tensors
- # by invoking build and compile on the model respectively, we do not have to
- # feed anything to the model. Model already has input and target data as
- # part of the graph.
- # Note: in this case, `any` and `all` are equivalent since we disallow
- # mixed symbolic/value inputs.
- if (not self.run_eagerly and is_build_called and is_compile_called and
- not from_iterator and any(_is_symbolic_tensor(v) for v in all_inputs)):
- return [], [], []
-
- # What follows is input validation and standardization to list format,
- # in the case where all inputs are value arrays.
-
- if self.run_eagerly:
- # In eager mode, do not do shape validation
- # since the network has no input nodes (placeholders) to be fed.
- feed_input_names = self.input_names
- feed_input_shapes = None
- elif not self._is_graph_network:
- # Case: symbolic-mode subclassed network. Do not do shape validation.
- feed_input_names = self._feed_input_names
- feed_input_shapes = None
- else:
- # Case: symbolic-mode graph network.
- # In this case, we run extensive shape validation checks.
- feed_input_names = self._feed_input_names
- feed_input_shapes = self._feed_input_shapes
-
- # Standardize the inputs.
- x = training_utils.standardize_input_data(
- x,
- feed_input_names,
- feed_input_shapes,
- check_batch_axis=False, # Don't enforce the batch size.
- exception_prefix='input')
-
- if y is not None:
- if not self._is_graph_network:
- feed_output_names = self._feed_output_names
- feed_output_shapes = None
- # Sample weighting not supported in this case.
- # TODO(fchollet): consider supporting it.
- feed_sample_weight_modes = [None for _ in self.outputs]
- else:
- feed_output_names = self._feed_output_names
- feed_sample_weight_modes = self._feed_sample_weight_modes
- feed_output_shapes = []
- for output_shape, loss_fn in zip(self._feed_output_shapes,
- self._feed_loss_fns):
- if loss_fn is losses.sparse_categorical_crossentropy:
- if K.image_data_format() == 'channels_first':
- feed_output_shapes.append(
- (output_shape[0], 1) + output_shape[2:])
- else:
- feed_output_shapes.append(output_shape[:-1] + (1,))
- elif (not hasattr(loss_fn, '__name__') or
- getattr(losses, loss_fn.__name__, None) is None):
- # If `loss_fn` is not a function (e.g. callable class)
- # or if it not in the `losses` module, then
- # it is a user-defined loss and we make no assumptions
- # about it.
- feed_output_shapes.append(None)
- else:
- feed_output_shapes.append(output_shape)
-
- # Standardize the outputs.
- y = training_utils.standardize_input_data(
- y,
- feed_output_names,
- # Don't enforce target shapes to match output shapes.
- # Precise checks will be run in `check_loss_and_target_compatibility`.
- shapes=None,
- check_batch_axis=False, # Don't enforce the batch size.
- exception_prefix='target')
-
- # Generate sample-wise weight values given the `sample_weight` and
- # `class_weight` arguments.
- sample_weights = training_utils.standardize_sample_weights(
- sample_weight, feed_output_names)
- class_weights = training_utils.standardize_class_weights(
- class_weight, feed_output_names)
- sample_weights = [
- training_utils.standardize_weights(ref, sw, cw, mode)
- for (ref, sw, cw, mode) in zip(y, sample_weights, class_weights,
- feed_sample_weight_modes)
- ]
- # Check that all arrays have the same length.
- if not self._distribution_strategy:
- training_utils.check_array_lengths(x, y, sample_weights)
- if self._is_graph_network and not self.run_eagerly:
- # Additional checks to avoid users mistakenly using improper loss fns.
- training_utils.check_loss_and_target_compatibility(
- y, self._feed_loss_fns, feed_output_shapes)
- else:
- y = []
- sample_weights = []
-
- if self.stateful and batch_size:
- # Check that for stateful networks, number of samples is a multiple
- # of the static batch size.
- if x[0].shape[0] % batch_size != 0:
- raise ValueError('In a stateful network, '
- 'you should only pass inputs with '
- 'a number of samples that can be '
- 'divided by the batch size. Found: ' +
- str(x[0].shape[0]) + ' samples')
-
- # If dictionary inputs were provided, we return a dictionary as well.
- if dict_inputs:
- x = dict(zip(feed_input_names, x))
- return x, y, sample_weights
-
- @checkpointable.no_automatic_dependency_tracking
- def _set_inputs(self, inputs, outputs=None, training=None):
- """Set model's input and output specs based on the input data received.
-
- This is to be used for Model subclasses, which do not know at instantiation
- time what their inputs look like.
-
- Args:
- inputs: Single array, or list of arrays. The arrays could be placeholders,
- Numpy arrays, data tensors, or TensorShapes.
- - if placeholders: the model is built on top of these placeholders,
- and we expect Numpy data to be fed for them when calling `fit`/etc.
- - if Numpy data or TensorShapes: we create placeholders matching the
- TensorShapes or shapes of the Numpy arrays. We expect Numpy data to be
- fed for these placeholders when calling `fit`/etc.
- - if data tensors: the model is built on top of these tensors.
- We do not expect any Numpy data to be provided when calling `fit`/etc.
- outputs: None, a data tensor, or a list of tensors. If None, the
- outputs will be determined by invoking `self.call()`, otherwise the
- provided value will be used.
- training: Boolean or None. Only relevant in symbolic mode. Specifies
- whether to build the model's graph in inference mode (False), training
- mode (True), or using the Keras learning phase (None).
- Raises:
- ValueError: If dict inputs are passed to a Sequential Model where the
- first layer isn't FeatureLayer.
- """
- if self.inputs:
- raise ValueError('Model inputs are already set.')
-
- if self.__class__.__name__ == 'Sequential' and not self.built:
- if tensor_util.is_tensor(inputs):
- input_shape = (None,) + tuple(inputs.shape.as_list()[1:])
- elif isinstance(inputs, tensor_shape.TensorShape):
- input_shape = (None,) + tuple(inputs.as_list()[1:])
- elif isinstance(inputs, dict):
- # We assert that the first layer is a FeatureLayer.
- if not training_utils.is_feature_layer(self.layers[0]):
- raise ValueError('Passing a dictionary input to a Sequential Model '
- 'which doesn\'t have FeatureLayer as the first layer'
- ' is an error.')
- input_shape = (None,)
- else:
- input_shape = (None,) + tuple(inputs.shape[1:])
- self._build_input_shape = input_shape
-
- # On-the-fly setting of symbolic model inputs (either by using the tensor
- # provided, or by creating a placeholder if Numpy data was provided).
- model_inputs = training_utils.ModelInputs(inputs)
- inputs = model_inputs.get_symbolic_inputs()
- self.inputs = model_inputs.get_symbolic_inputs(return_single_as_list=True)
- self.input_names = model_inputs.get_input_names()
-
- self._feed_inputs = []
- self._feed_input_names = []
- self._feed_input_shapes = []
-
- for k, v in model_inputs.as_dict():
- if K.is_placeholder(v):
- self._feed_inputs.append(v)
- self._feed_input_names.append(k)
- self._feed_input_shapes.append(K.int_shape(v))
-
- # TODO(fchollet): consider calling `_maybe_build` before calling the model.
-
- if outputs is None:
- # Obtain symbolic outputs by calling the model.
- with K.get_graph().as_default():
- if self._expects_training_arg:
- outputs = self.call(inputs, training=training)
- else:
- outputs = self.call(inputs)
-
- outputs = nest.flatten(outputs)
- self.outputs = outputs
- self.output_names = [
- 'output_%d' % (i + 1) for i in range(len(self.outputs))]
- self.built = True
+ @run_eagerly.setter
+ def run_eagerly(self, value):
+ self._run_eagerly = value
def fit(self,
x=None,
@@ -1608,7 +630,7 @@
0 = silent, 1 = progress bar, 2 = one line per epoch.
callbacks: List of `keras.callbacks.Callback` instances.
List of callbacks to apply during training.
- See [callbacks](/api_docs/python/tf/keras/callbacks).
+ See `tf.keras.callbacks`.
validation_split: Float between 0 and 1.
Fraction of the training data to be used as validation data.
The model will set apart this fraction of the training data,
@@ -2095,6 +1117,8 @@
if hasattr(self, 'metrics'):
for m in self.metrics:
m.reset_states()
+ if self._distribution_strategy:
+ training_distributed._reset_metrics(self) # pylint: disable=protected-access
def train_on_batch(self,
x,
@@ -2584,6 +1608,933 @@
batch_size = 32
return batch_size
+ @property
+ def _default_save_signature(self):
+ return training_utils.trace_model_call(self)
+
+ def _set_sample_weight_attributes(self, sample_weight_mode,
+ skip_target_weighing_indices):
+ """Sets sample weight related attributes on the model."""
+ sample_weights, sample_weight_modes = training_utils.prepare_sample_weights(
+ self.output_names, sample_weight_mode, skip_target_weighing_indices)
+ self.sample_weights = sample_weights
+ self.sample_weight_modes = sample_weight_modes
+ self._feed_sample_weight_modes = [
+ sample_weight_modes[i]
+ for i in range(len(self.outputs))
+ if i not in skip_target_weighing_indices
+ ]
+ self._feed_sample_weights = [
+ sample_weights[i]
+ for i in range(len(sample_weights))
+ if i not in skip_target_weighing_indices
+ ]
+
+ def _cache_output_metric_attributes(self, metrics, weighted_metrics):
+ """Caches metric name and function attributes for every model output."""
+ output_shapes = [
+ None if output is None else output.get_shape().as_list()
+ for output in self.outputs
+ ]
+ self._per_output_metrics = training_utils.collect_per_output_metric_info(
+ metrics, self.output_names, output_shapes, self.loss_functions)
+ self._per_output_weighted_metrics = \
+ training_utils.collect_per_output_metric_info(
+ weighted_metrics, self.output_names, output_shapes,
+ self.loss_functions, self.sample_weights)
+
+ def _add_unique_metric_name(self, metric_name, output_index):
+ """Makes the metric name unique and adds it to the model's metric name list.
+
+ If there are multiple outputs for which the metrics are calculated, the
+ metric names have to be made unique by appending an integer.
+
+ Arguments:
+ metric_name: Metric name that corresponds to the metric specified by the
+ user. For example: 'acc'.
+ output_index: The index of the model output for which the metric name is
+ being added.
+
+ Returns:
+ string, name of the model's unique metric name
+ """
+ if len(self.output_names) > 1:
+ metric_name = '%s_%s' % (self.output_names[output_index], metric_name)
+ j = 1
+ base_metric_name = metric_name
+ while metric_name in self._compile_metrics_names:
+ metric_name = '%s_%d' % (base_metric_name, j)
+ j += 1
+
+ return metric_name
+
+ @property
+ def _all_metrics_tensors(self):
+ """Returns the network's symbolic metric tensors."""
+ metrics_tensors = {}
+ if self._is_compiled:
+ metrics_tensors.update(self._compile_metrics_tensors)
+ metrics_tensors.update(super(Model, self)._all_metrics_tensors)
+ return metrics_tensors
+
+ @property
+ def _all_stateful_metrics_tensors(self):
+ """Returns the network's symbolic metric tensors."""
+ metrics_tensors = {}
+ if self._is_compiled:
+ metrics_tensors.update(self._compile_stateful_metrics_tensors)
+ metrics_tensors.update(super(Model, self)._all_metrics_tensors)
+ return metrics_tensors
+
+ def _init_metric_attributes(self):
+ """Initialized model metric attributes."""
+ # List of all metric names in the model.
+ self._compile_metrics_names = ['loss']
+ # List of stateful metric functions. Used for resetting metric state during
+ # training/eval.
+ # This includes loss functions when there are multiple outputs.
+ self._compile_stateful_metric_functions = []
+ # Dict of all aggregated metric result tensors. This includes aggregated
+ # loss result tensors when there are multiple outputs.
+ self._compile_stateful_metrics_tensors = {}
+ # Dict of all metric result tensors (aggregated or not - based on the
+ # values given in compile.). This includes aggregated loss result tensors
+ # when there are multiple outputs.
+ self._compile_metrics_tensors = {}
+
+ def _set_per_output_metric_attributes(self, metrics_dict, output_index):
+ """Sets the metric attributes on the model for the given output.
+
+ Arguments:
+ metrics_dict: A dict with metric names as keys and metric fns as values.
+ output_index: The index of the model output for which the metric
+ attributes are added.
+
+ Returns:
+ Metrics dict updated with unique metric names as keys.
+ """
+ updated_metrics_dict = collections.OrderedDict()
+ for metric_name, (metric_fn, stateful_metric_fn) in metrics_dict.items():
+ metric_name = self._add_unique_metric_name(metric_name, output_index)
+ updated_metrics_dict[metric_name] = (metric_fn, stateful_metric_fn)
+ # Keep track of metric name, function and stateful function.
+ self._compile_metrics_names.append(metric_name)
+ self._compile_stateful_metric_functions.append(stateful_metric_fn)
+ return updated_metrics_dict
+
+ def _set_metric_attributes(self, outputs, skip_target_indices=None):
+ """Sets the metric attributes on the model for all the model outputs."""
+ skip_target_indices = skip_target_indices or []
+ updated_per_output_metrics = []
+ updated_per_output_weighted_metrics = []
+ for i in range(len(outputs)):
+ if i in skip_target_indices:
+ updated_per_output_metrics.append(self._per_output_metrics[i])
+ updated_per_output_weighted_metrics.append(
+ self._per_output_weighted_metrics[i])
+ continue
+ updated_per_output_metrics.append(
+ self._set_per_output_metric_attributes(self._per_output_metrics[i],
+ i))
+ updated_per_output_weighted_metrics.append(
+ self._set_per_output_metric_attributes(
+ self._per_output_weighted_metrics[i], i))
+
+ self._per_output_metrics = updated_per_output_metrics
+ self._per_output_weighted_metrics = updated_per_output_weighted_metrics
+
+ def _handle_per_output_metrics(self,
+ metrics_dict,
+ y_true,
+ y_pred,
+ mask,
+ weights=None,
+ return_stateful_result=True):
+ """Calls metric functions for a single output.
+
+ Arguments:
+ metrics_dict: A dict with metric names as keys and metric fns as values.
+ y_true: Target output.
+ y_pred: Predicted output.
+ mask: Computed mask value for the current output.
+ weights: Weights to be applied on the current output.
+ return_stateful_result: Boolean, indicates whether the stateful
+ (aggregated)/stateless metric result should be returned.
+
+ Returns:
+ A list of metric result tensors.
+ """
+ metric_results = []
+ for metric_name, (metric_fn, stateful_fn) in metrics_dict.items():
+ with K.name_scope(metric_name):
+
+ def _call_stateful_fn(fn):
+ return training_utils.call_metric_function(
+ fn, y_true, y_pred, weights=weights, mask=mask)
+
+ def _call_stateless_fn(fn):
+ weighted_metric_fn = training_utils.weighted_masked_objective(fn)
+ return weighted_metric_fn(y_true, y_pred, weights=weights, mask=mask)
+
+ def _track_metric_tensors(name, stateless_result, stateful_result):
+ self._compile_metrics_tensors[name] = stateless_result
+ self._compile_stateful_metrics_tensors[name] = stateful_result
+
+ if isinstance(metric_fn, metrics_module.Metric):
+ # If the given metric fn is stateful, call the fn and return result.
+ metric_result = _call_stateful_fn(metric_fn)
+ metric_results.append(metric_result)
+ if not self.run_eagerly:
+ _track_metric_tensors(metric_name, metric_result, metric_result)
+ elif self.run_eagerly:
+ # In eager mode, if the given metric fn is not stateful, we invoke the
+ # given fn or its stateful version based on the given flag.
+ if return_stateful_result:
+ metric_result = _call_stateful_fn(stateful_fn)
+ else:
+ metric_result = _call_stateless_fn(metric_fn)
+ metric_results.append(metric_result)
+ else:
+ # In graph mode, we build the sub-graph for both the stateful and the
+ # stateless fns.
+ stateful_metric_result = _call_stateful_fn(stateful_fn)
+ metric_result = _call_stateless_fn(metric_fn)
+ _track_metric_tensors(metric_name, metric_result,
+ stateful_metric_result)
+
+ return metric_results
+
+ def _handle_metrics(self,
+ outputs,
+ skip_target_indices=None,
+ targets=None,
+ sample_weights=None,
+ masks=None,
+ return_stateful_result=True):
+ """Handles calling metric functions.
+
+ Arguments:
+ outputs: List of outputs (predictions).
+ skip_target_indices: Optional. List of target ids to skip.
+ targets: List of targets.
+ sample_weights: Optional list of sample weight arrays.
+ masks: List of computed output mask values.
+ return_stateful_result: Boolean, indicates whether the stateful
+ (aggregated)/stateless metric result should be returned.
+
+ Returns:
+ A list of metric result tensors.
+ """
+ skip_target_indices = skip_target_indices or []
+ metric_results = []
+ with K.name_scope('metrics'):
+ # Invoke all metrics added using `compile`.
+ for i in range(len(outputs)):
+ if i in skip_target_indices:
+ continue
+ output = outputs[i] if outputs else None
+ target = targets[i] if targets else None
+ output_mask = masks[i] if masks else None
+ metric_results.extend(
+ self._handle_per_output_metrics(
+ self._per_output_metrics[i],
+ target,
+ output,
+ output_mask,
+ return_stateful_result=return_stateful_result))
+ metric_results.extend(
+ self._handle_per_output_metrics(
+ self._per_output_weighted_metrics[i],
+ target,
+ output,
+ output_mask,
+ weights=sample_weights[i],
+ return_stateful_result=return_stateful_result))
+
+ # Add metric results from the `add_metric` metrics in eager mode.
+ if context.executing_eagerly():
+ for m in self.metrics:
+ if m not in self._compile_stateful_metric_functions:
+ metric_results.append(m.result())
+ return metric_results
+
+ def _check_trainable_weights_consistency(self):
+ """Check trainable weights count consistency.
+
+ This will raise a warning if `trainable_weights` and
+ `_collected_trainable_weights` are inconsistent (i.e. have different
+ number of parameters).
+ Inconsistency will typically arise when one modifies `model.trainable`
+ without calling `model.compile` again.
+ """
+ if not hasattr(self, '_collected_trainable_weights'):
+ return
+
+ if len(self.trainable_weights) != len(self._collected_trainable_weights):
+ logging.log_first_n(
+ logging.WARN, 'Discrepancy between trainable weights and collected'
+ ' trainable weights, did you set `model.trainable`'
+ ' without calling `model.compile` after ?', 1)
+
+ def _make_train_function_helper(self, fn_name, outputs, metric_updates=None):
+ if not hasattr(self, fn_name):
+ raise RuntimeError('You must compile your model before using it.')
+ self._check_trainable_weights_consistency()
+ if getattr(self, fn_name) is None:
+ inputs = (self._feed_inputs +
+ self._feed_targets +
+ self._feed_sample_weights)
+ if not isinstance(K.symbolic_learning_phase(), int):
+ inputs += [K.symbolic_learning_phase()]
+
+ with K.get_graph().as_default():
+ with K.name_scope('training'):
+ with K.name_scope(self.optimizer.__class__.__name__):
+ # Training updates
+ updates = self.optimizer.get_updates(
+ params=self._collected_trainable_weights, loss=self.total_loss)
+ # Unconditional updates
+ updates += self.get_updates_for(None)
+ # Conditional updates relevant to this model
+ updates += self.get_updates_for(self.inputs)
+ # Add stateful metrics updates.
+ if metric_updates is not None:
+ updates += metric_updates
+
+ with K.name_scope('training'):
+ # Gets loss and metrics. Updates weights at each call.
+ fn = K.function(
+ inputs,
+ outputs,
+ updates=updates,
+ name='train_function',
+ **self._function_kwargs)
+ setattr(self, fn_name, fn)
+
+ def _make_train_function(self):
+ metrics_tensors = [
+ self._all_metrics_tensors[m] for m in self.metrics_names[1:]
+ ]
+ self._make_train_function_helper('train_function',
+ [self.total_loss] + metrics_tensors)
+
+ def _make_fit_function(self):
+ metrics_tensors = [
+ self._all_stateful_metrics_tensors[m] for m in self.metrics_names[1:]
+ ]
+ self._make_train_function_helper(
+ '_fit_function', [self.total_loss] + metrics_tensors)
+
+ def _make_test_function_helper(self, fn_name, outputs, metric_updates=None):
+ if not hasattr(self, fn_name):
+ raise RuntimeError('You must compile your model before using it.')
+ if getattr(self, fn_name) is None:
+ inputs = (self._feed_inputs +
+ self._feed_targets +
+ self._feed_sample_weights)
+
+ with K.name_scope('evaluation'):
+ updates = self.state_updates
+ # Add stateful metrics updates.
+ if metric_updates is not None:
+ updates += metric_updates
+ # Return loss and metrics, no gradient updates.
+ # Does update the network states.
+ fn = K.function(
+ inputs,
+ outputs,
+ updates=updates,
+ name='test_function',
+ **self._function_kwargs)
+ setattr(self, fn_name, fn)
+
+ def _make_test_function(self):
+ metrics_tensors = [
+ self._all_metrics_tensors[m] for m in self.metrics_names[1:]
+ ]
+ self._make_test_function_helper('test_function',
+ [self.total_loss] + metrics_tensors)
+
+ def _make_eval_function(self):
+ metrics_tensors = [
+ self._all_stateful_metrics_tensors[m] for m in self.metrics_names[1:]
+ ]
+ self._make_test_function_helper(
+ '_eval_function', [self.total_loss] + metrics_tensors)
+
+ def _make_predict_function(self):
+ if not hasattr(self, 'predict_function'):
+ self.predict_function = None
+ if self.predict_function is None:
+ inputs = self._feed_inputs
+ # Gets network outputs. Does not update weights.
+ # Does update the network states.
+ kwargs = getattr(self, '_function_kwargs', {})
+ with K.name_scope('predict'):
+ self.predict_function = K.function(
+ inputs,
+ self.outputs,
+ updates=self.state_updates,
+ name='predict_function',
+ **kwargs)
+
+ def _make_execution_function(self, mode):
+ if mode == 'train':
+ self._make_fit_function()
+ return self._fit_function
+ if mode == 'test':
+ self._make_eval_function()
+ return self._eval_function
+ if mode == 'predict':
+ self._make_predict_function()
+ return self.predict_function
+
+ def _get_iterator_get_next_tensors(self, iterator):
+ get_next_op = self._iterator_get_next.get(iterator, None)
+ if get_next_op is None:
+ get_next_op = iterator.get_next()
+ self._iterator_get_next[iterator] = get_next_op
+ return get_next_op
+
+ def _distribution_standardize_user_data(self,
+ x,
+ y=None,
+ sample_weight=None,
+ class_weight=None,
+ batch_size=None,
+ check_steps=False,
+ steps_name='steps',
+ steps=None,
+ validation_split=0,
+ shuffle=False):
+ """Runs validation checks on input and target data passed by the user.
+
+ This is called when using DistributionStrategy to train, evaluate or serve
+ the model.
+
+ Args:
+ x: Input data. A numpy array or `tf.data` dataset.
+ y: Target data. A numpy array or None if x is a `tf.data` dataset.
+ sample_weight: An optional sample-weight array passed by the user to
+ weight the importance of each sample in `x`.
+ class_weight: An optional class-weight array by the user to
+ weight the importance of samples in `x` based on the class they belong
+ to, as conveyed by `y`.
+ batch_size: Integer batch size. If provided, it is used to run additional
+ validation checks on stateful models.
+ check_steps: boolean, True if we want to check for validity of `steps` and
+ False, otherwise.
+ steps_name: The public API's parameter name for `steps`.
+ steps: Integer or `None`. Total number of steps (batches of samples) to
+ execute.
+ validation_split: Float between 0 and 1.
+ Fraction of the training data to be used as validation data.
+ shuffle: Boolean whether to shuffle the training data before each epoch.
+
+ Returns:
+ Iterator for reading the dataset `x`.
+
+ Raises:
+ ValueError: In case of invalid user-provided data.
+ RuntimeError: If the model was never compiled.
+ """
+ if class_weight:
+ raise NotImplementedError('`class_weight` is currently not supported '
+ 'when using DistributionStrategy.')
+
+ if (sample_weight is not None and sample_weight.all() and
+ distributed_training_utils.is_tpu_strategy(
+ self._distribution_strategy)):
+ raise NotImplementedError('`sample_weight` is currently not supported '
+ 'when using TPUStrategy.')
+
+ # Validates `steps` argument right at the beginning since we use it to
+ # construct the dataset object.
+ # TODO(anjalisridhar): Remove this check once we refactor the
+ # _standardize_user_data code path. This check is already present elsewhere
+ # in the codebase.
+ if check_steps and isinstance(x, dataset_ops.DatasetV2) and steps is None:
+ raise ValueError('When using Datasets as input, '
+ 'you should specify the `{steps_name}` argument.'
+ .format(steps_name=steps_name))
+
+ first_x_value = nest.flatten(x)[0]
+ if isinstance(first_x_value, np.ndarray):
+ # We need to use the drop_remainder argument to allow for a static
+ # input shape which is required for TPUs.
+ drop_remainder = self._distribution_strategy.require_static_shapes
+ if y is not None:
+ var_x = distributed_training_utils.get_var_for_numpy(
+ self._distribution_strategy, x)
+ var_y = distributed_training_utils.get_var_for_numpy(
+ self._distribution_strategy, y)
+ if sample_weight is not None:
+ var_sample_weights = distributed_training_utils.get_var_for_numpy(
+ self._distribution_strategy, sample_weight)
+
+ x = dataset_ops.Dataset.from_tensor_slices((var_x, var_y,
+ var_sample_weights))
+ else:
+ x = dataset_ops.Dataset.from_tensor_slices((var_x, var_y))
+
+ x = dataset_ops.Dataset.from_tensor_slices((var_x, var_y))
+ if shuffle:
+ # 1024 is a good buffer size since it is much larger than the average
+ # batch size provided by the user and provides sufficient randomness.
+ # One thing to keep in mind is the memory usage based on the size of
+ # each sample.
+ x = x.shuffle(1024)
+ x = x.repeat()
+ x = x.batch(batch_size, drop_remainder=drop_remainder)
+ y = None
+ sample_weight = None
+ else:
+ # This case is for the predict call where the dataset only contains
+ # inputs and no targets, i.e. it does not return a tuple
+ var_x = distributed_training_utils.get_var_for_numpy(
+ self._distribution_strategy, x)
+ x = dataset_ops.Dataset.from_tensor_slices(var_x)
+ x = x.batch(batch_size, drop_remainder=drop_remainder)
+
+ assert isinstance(x, dataset_ops.DatasetV2)
+
+ with self._distribution_strategy.scope():
+ iterator = self._distribution_strategy.make_dataset_iterator(x)
+ init_op = iterator.initialize()
+ if not context.executing_eagerly():
+ K.get_session().run(init_op)
+
+ training_utils.validate_iterator_input(x, y, sample_weight,
+ validation_split)
+ return iterator
+
+ def _standardize_user_data(self,
+ x,
+ y=None,
+ sample_weight=None,
+ class_weight=None,
+ batch_size=None,
+ check_steps=False,
+ steps_name='steps',
+ steps=None,
+ validation_split=0,
+ shuffle=False):
+ """Runs validation checks on input and target data passed by the user.
+
+ Also standardizes the data to lists of arrays, in order.
+
+ Also builds and compiles the model on the fly if it is a subclassed model
+ that has never been called before (and thus has no inputs/outputs).
+
+ This is a purely internal method, subject to refactoring at any time.
+
+ Args:
+ x: Input data. It could be:
+ - A Numpy array (or array-like), or a list of arrays
+ (in case the model has multiple inputs).
+ - A TensorFlow tensor, or a list of tensors
+ (in case the model has multiple inputs).
+ - A dict mapping input names to the corresponding array/tensors,
+ if the model has named inputs.
+ - A `tf.data` dataset or a dataset iterator.
+ y: Target data. Like the input data `x`,
+ it could be either Numpy array(s) or TensorFlow tensor(s).
+ It should be consistent with `x` (you cannot have Numpy inputs and
+ tensor targets, or inversely). If `x` is a dataset or a
+ dataset iterator, `y` should not be specified
+ (since targets will be obtained from the iterator).
+ sample_weight: An optional sample-weight array passed by the user to
+ weight the importance of each sample in `x`.
+ class_weight: An optional class-weight array by the user to
+ weight the importance of samples in `x` based on the class they belong
+ to, as conveyed by `y`.
+ batch_size: Integer batch size. If provided, it is used to run additional
+ validation checks on stateful models.
+ check_steps: boolean, True if we want to check for validity of `steps` and
+ False, otherwise. For example, when we are standardizing one batch of
+ data for train_on_batch/predict_on_batch/test_on_batch APIs, `steps`
+ value is not required and we should not check for its validity in these
+ cases.
+ steps_name: The public API's parameter name for `steps`.
+ steps: Integer or `None`. Total number of steps (batches of samples) to
+ execute.
+ validation_split: Float between 0 and 1.
+ Fraction of the training data to be used as validation data.
+ shuffle: Boolean whether to shuffle the training data before each epoch.
+
+ Returns:
+ A tuple of 3: inputs (arrays or dicts, depending on whether `x` was a dict
+ or not), target arrays, sample-weight arrays.
+ If the model's input and targets are symbolic, these lists are empty
+ (since the model takes no user-provided data, instead the data comes
+ from the symbolic inputs/targets).
+
+ Raises:
+ ValueError: In case of invalid user-provided data.
+ RuntimeError: If the model was never compiled.
+ """
+ if self._distribution_strategy:
+ iterator = self._distribution_standardize_user_data(
+ x,
+ y,
+ sample_weight=sample_weight,
+ class_weight=class_weight,
+ batch_size=batch_size,
+ check_steps=check_steps,
+ steps_name=steps_name,
+ steps=steps,
+ validation_split=validation_split,
+ shuffle=shuffle)
+ return iterator, None, None
+
+ if isinstance(x, dataset_ops.DatasetV2):
+ if context.executing_eagerly():
+ x = iter(x)
+ else:
+ if x in self._dataset_iterator_cache:
+ x = self._dataset_iterator_cache[x]
+ else:
+ iterator = dataset_ops.make_initializable_iterator(x)
+ self._dataset_iterator_cache[x] = iterator
+ x = iterator
+ K.get_session().run(x.initializer)
+
+ # Validates `steps` argument based on x's type.
+ if check_steps:
+ training_utils.check_steps_argument(x, steps, steps_name)
+
+ is_x_eager_iterator = isinstance(x, iterator_ops.EagerIterator)
+ is_x_iterator = isinstance(x, iterator_ops.Iterator)
+
+ # Validate user inputs when data is given as a dataset or dataset iterator.
+ if is_x_iterator or is_x_eager_iterator:
+ training_utils.validate_iterator_input(x, y, sample_weight,
+ validation_split)
+
+ # For eager iterators, when we have to process multiple batches of samples,
+ # we will standardize the data when we actually loop over iterator and get
+ # the batches. For now, we just return the iterator as is.
+ if is_x_eager_iterator:
+ return x, y, sample_weight
+
+ # If input data is a dataset iterator in graph mode or if it is an eager
+ # iterator and only one batch of samples is required, we fetch the data
+ # tensors from the iterator and then standardize them.
+ if is_x_iterator or is_x_eager_iterator:
+ try:
+ if is_x_iterator:
+ next_element = self._get_iterator_get_next_tensors(x)
+ else:
+ next_element = x.get_next()
+ except errors.OutOfRangeError:
+ raise RuntimeError('Your dataset iterator ran out of data; '
+ 'Make sure that your dataset can generate '
+ 'required number of samples.')
+
+ if isinstance(next_element, (list, tuple)):
+ if len(next_element) not in [2, 3]:
+ raise ValueError(
+ 'Please provide model inputs as a list or tuple of 2 or 3'
+ 'elements: (input, target) or (input, target, sample_weights)'
+ 'Received %s' % next_element)
+ if len(next_element) == 2:
+ x, y = next_element
+ else:
+ x, y, sample_weight = next_element
+ else:
+ x = next_element
+
+ if sample_weight is not None and class_weight is not None:
+ logging.warning(
+ 'Received both a `sample_weight` and `class_weight` argument. '
+ 'The `class_weight` argument will be ignored.')
+ # First, we build/compile the model on the fly if necessary.
+ all_inputs = []
+ is_build_called = False
+ is_compile_called = False
+ # Whether this is a subclassed model that expects dictionary inputs
+ # rather than list inputs (e.g. FeatureColumn-based models).
+ dict_inputs = False
+ if not self.inputs:
+ # We need to use `x` to set the model inputs.
+ # We type-check that `x` and `y` are either single arrays
+ # or lists of arrays.
+ if isinstance(x, (list, tuple)):
+ if not all(isinstance(v, np.ndarray) or
+ tensor_util.is_tensor(v) for v in x):
+ raise ValueError('Please provide as model inputs either a single '
+ 'array or a list of arrays. You passed: x=' + str(x))
+ all_inputs += list(x)
+ elif isinstance(x, dict):
+ dict_inputs = True
+ keys = sorted(x.keys())
+ all_inputs = [x[k] for k in keys]
+ else:
+ if not isinstance(x, np.ndarray) and not tensor_util.is_tensor(x):
+ raise ValueError('Please provide as model inputs either a single '
+ 'array or a list of arrays. You passed: x=' + str(x))
+ all_inputs.append(x)
+
+ # Build the model using the retrieved inputs (value or symbolic).
+ # If values or generated from a dataset, then in symbolic-mode
+ # placeholders will be created to match the value shapes.
+ if not self.inputs:
+ is_build_called = True
+ if is_x_iterator:
+ cast_inputs = nest.map_structure(lambda v: v.shape, x)
+ elif training_utils.has_tensors(x):
+ cast_inputs = training_utils.cast_if_floating_dtype(x)
+ else:
+ cast_inputs = x
+ self._set_inputs(cast_inputs)
+ else:
+ dict_inputs = isinstance(self.inputs, dict)
+ if dict_inputs and context.executing_eagerly():
+ # No support for graph functions when the model expects dictionary inputs
+ # (i.e. FeatureColumn-based models).
+ self.run_eagerly = True
+
+ if y is not None:
+ if not self.optimizer:
+ raise RuntimeError('You must compile a model before '
+ 'training/testing. '
+ 'Use `model.compile(optimizer, loss)`.')
+ if not self._is_compiled:
+ # On-the-fly compilation of the model.
+ # We need to use `y` to set the model targets.
+ if training_utils.has_tensors(y):
+ y = training_utils.cast_if_floating_dtype(y)
+ if isinstance(y, (list, tuple)):
+ if not all(isinstance(v, np.ndarray) or
+ tensor_util.is_tensor(v) for v in y):
+ raise ValueError('Please provide as model targets either a single '
+ 'array or a list of arrays. '
+ 'You passed: y=' + str(y))
+ all_inputs += list(y)
+ elif isinstance(y, dict):
+ raise ValueError('Please do not pass a dictionary as model targets.')
+ else:
+ if not isinstance(y, np.ndarray) and not tensor_util.is_tensor(y):
+ raise ValueError('Please provide as model targets either a single '
+ 'array or a list of arrays. '
+ 'You passed: y=' + str(y))
+ all_inputs.append(y)
+
+ # Typecheck that all inputs are *either* value *or* symbolic.
+ # TODO(fchollet): this check could be removed in Eager mode?
+ if any(tensor_util.is_tensor(v) for v in all_inputs):
+ if not all(tensor_util.is_tensor(v) for v in all_inputs):
+ raise ValueError('Do not pass inputs that mix Numpy arrays and '
+ 'TensorFlow tensors. '
+ 'You passed: x=' + str(x) + '; y=' + str(y))
+
+ if self.run_eagerly or is_x_iterator:
+ target_tensors = None
+ else:
+ # Handle target tensors if any passed.
+ if not isinstance(y, (list, tuple)):
+ y = [y]
+ target_tensors = [v for v in y if _is_symbolic_tensor(v)]
+ is_compile_called = True
+ self.compile(
+ optimizer=self.optimizer,
+ loss=self.loss,
+ metrics=self._compile_metrics,
+ weighted_metrics=self._compile_weighted_metrics,
+ loss_weights=self.loss_weights,
+ target_tensors=target_tensors,
+ run_eagerly=self.run_eagerly)
+
+ # In graph mode, if we had just set inputs and targets as symbolic tensors
+ # by invoking build and compile on the model respectively, we do not have to
+ # feed anything to the model. Model already has input and target data as
+ # part of the graph.
+ # Note: in this case, `any` and `all` are equivalent since we disallow
+ # mixed symbolic/value inputs.
+ if (not self.run_eagerly and is_build_called and is_compile_called and
+ not is_x_iterator and any(_is_symbolic_tensor(v) for v in all_inputs)):
+ return [], [], []
+
+ # What follows is input validation and standardization to list format,
+ # in the case where all inputs are value arrays.
+
+ if self.run_eagerly:
+ # In eager mode, do not do shape validation
+ # since the network has no input nodes (placeholders) to be fed.
+ feed_input_names = self.input_names
+ feed_input_shapes = None
+ elif not self._is_graph_network:
+ # Case: symbolic-mode subclassed network. Do not do shape validation.
+ feed_input_names = self._feed_input_names
+ feed_input_shapes = None
+ else:
+ # Case: symbolic-mode graph network.
+ # In this case, we run extensive shape validation checks.
+ feed_input_names = self._feed_input_names
+ feed_input_shapes = self._feed_input_shapes
+
+ # Standardize the inputs.
+ x = training_utils.standardize_input_data(
+ x,
+ feed_input_names,
+ feed_input_shapes,
+ check_batch_axis=False, # Don't enforce the batch size.
+ exception_prefix='input')
+
+ if y is not None:
+ if not self._is_graph_network:
+ feed_output_names = self._feed_output_names
+ feed_output_shapes = None
+ # Sample weighting not supported in this case.
+ # TODO(fchollet): consider supporting it.
+ feed_sample_weight_modes = [None for _ in self.outputs]
+ else:
+ feed_output_names = self._feed_output_names
+ feed_sample_weight_modes = self._feed_sample_weight_modes
+ feed_output_shapes = []
+ for output_shape, loss_fn in zip(self._feed_output_shapes,
+ self._feed_loss_fns):
+ if loss_fn is losses.sparse_categorical_crossentropy:
+ if K.image_data_format() == 'channels_first':
+ feed_output_shapes.append(
+ (output_shape[0], 1) + output_shape[2:])
+ else:
+ feed_output_shapes.append(output_shape[:-1] + (1,))
+ elif (not hasattr(loss_fn, '__name__') or
+ getattr(losses, loss_fn.__name__, None) is None):
+ # If `loss_fn` is not a function (e.g. callable class)
+ # or if it not in the `losses` module, then
+ # it is a user-defined loss and we make no assumptions
+ # about it.
+ feed_output_shapes.append(None)
+ else:
+ feed_output_shapes.append(output_shape)
+
+ # Standardize the outputs.
+ y = training_utils.standardize_input_data(
+ y,
+ feed_output_names,
+ # Don't enforce target shapes to match output shapes.
+ # Precise checks will be run in `check_loss_and_target_compatibility`.
+ shapes=None,
+ check_batch_axis=False, # Don't enforce the batch size.
+ exception_prefix='target')
+
+ # Generate sample-wise weight values given the `sample_weight` and
+ # `class_weight` arguments.
+ sample_weights = training_utils.standardize_sample_weights(
+ sample_weight, feed_output_names)
+ class_weights = training_utils.standardize_class_weights(
+ class_weight, feed_output_names)
+ sample_weights = [
+ training_utils.standardize_weights(ref, sw, cw, mode)
+ for (ref, sw, cw, mode) in zip(y, sample_weights, class_weights,
+ feed_sample_weight_modes)
+ ]
+ # Check that all arrays have the same length.
+ if not self._distribution_strategy:
+ training_utils.check_array_lengths(x, y, sample_weights)
+ if self._is_graph_network and not self.run_eagerly:
+ # Additional checks to avoid users mistakenly using improper loss fns.
+ training_utils.check_loss_and_target_compatibility(
+ y, self._feed_loss_fns, feed_output_shapes)
+ else:
+ y = []
+ sample_weights = []
+
+ if self.stateful and batch_size:
+ # Check that for stateful networks, number of samples is a multiple
+ # of the static batch size.
+ if x[0].shape[0] % batch_size != 0:
+ raise ValueError('In a stateful network, '
+ 'you should only pass inputs with '
+ 'a number of samples that can be '
+ 'divided by the batch size. Found: ' +
+ str(x[0].shape[0]) + ' samples')
+
+ # If dictionary inputs were provided, we return a dictionary as well.
+ if dict_inputs:
+ x = dict(zip(feed_input_names, x))
+ return x, y, sample_weights
+
+ @checkpointable.no_automatic_dependency_tracking
+ def _set_inputs(self, inputs, outputs=None, training=None):
+ """Set model's input and output specs based on the input data received.
+
+ This is to be used for Model subclasses, which do not know at instantiation
+ time what their inputs look like.
+
+ Args:
+ inputs: Single array, or list of arrays. The arrays could be placeholders,
+ Numpy arrays, data tensors, or TensorShapes.
+ - if placeholders: the model is built on top of these placeholders,
+ and we expect Numpy data to be fed for them when calling `fit`/etc.
+ - if Numpy data or TensorShapes: we create placeholders matching the
+ TensorShapes or shapes of the Numpy arrays. We expect Numpy data to be
+ fed for these placeholders when calling `fit`/etc.
+ - if data tensors: the model is built on top of these tensors.
+ We do not expect any Numpy data to be provided when calling `fit`/etc.
+ outputs: None, a data tensor, or a list of tensors. If None, the
+ outputs will be determined by invoking `self.call()`, otherwise the
+ provided value will be used.
+ training: Boolean or None. Only relevant in symbolic mode. Specifies
+ whether to build the model's graph in inference mode (False), training
+ mode (True), or using the Keras learning phase (None).
+ Raises:
+ ValueError: If dict inputs are passed to a Sequential Model where the
+ first layer isn't FeatureLayer.
+ """
+ if self.inputs:
+ raise ValueError('Model inputs are already set.')
+
+ if self.__class__.__name__ == 'Sequential' and not self.built:
+ if tensor_util.is_tensor(inputs):
+ input_shape = (None,) + tuple(inputs.shape.as_list()[1:])
+ elif isinstance(inputs, tensor_shape.TensorShape):
+ input_shape = (None,) + tuple(inputs.as_list()[1:])
+ elif isinstance(inputs, dict):
+ # We assert that the first layer is a FeatureLayer.
+ if not training_utils.is_feature_layer(self.layers[0]):
+ raise ValueError('Passing a dictionary input to a Sequential Model '
+ 'which doesn\'t have FeatureLayer as the first layer'
+ ' is an error.')
+ input_shape = (None,)
+ else:
+ input_shape = (None,) + tuple(inputs.shape[1:])
+ self._build_input_shape = input_shape
+
+ # On-the-fly setting of symbolic model inputs (either by using the tensor
+ # provided, or by creating a placeholder if Numpy data was provided).
+ model_inputs = training_utils.ModelInputs(inputs)
+ inputs = model_inputs.get_symbolic_inputs()
+ self.inputs = model_inputs.get_symbolic_inputs(return_single_as_list=True)
+ self.input_names = model_inputs.get_input_names()
+
+ self._feed_inputs = []
+ self._feed_input_names = []
+ self._feed_input_shapes = []
+
+ for k, v in model_inputs.as_dict():
+ if K.is_placeholder(v):
+ self._feed_inputs.append(v)
+ self._feed_input_names.append(k)
+ self._feed_input_shapes.append(K.int_shape(v))
+
+ # TODO(fchollet): consider calling `_maybe_build` before calling the model.
+
+ if outputs is None:
+ # Obtain symbolic outputs by calling the model.
+ with K.get_graph().as_default():
+ if self._expects_training_arg:
+ outputs = self.call(inputs, training=training)
+ else:
+ outputs = self.call(inputs)
+
+ outputs = nest.flatten(outputs)
+ self.outputs = outputs
+ self.output_names = training_utils.generic_output_names(outputs)
+ self.built = True
+
class DistributedCallbackModel(Model):
"""Model that is used for callbacks with DistributionStrategy."""
diff --git a/tensorflow/python/keras/engine/training_arrays.py b/tensorflow/python/keras/engine/training_arrays.py
index e9dfbcb..196d48f 100644
--- a/tensorflow/python/keras/engine/training_arrays.py
+++ b/tensorflow/python/keras/engine/training_arrays.py
@@ -138,6 +138,7 @@
steps_per_epoch=None,
validation_steps=None,
mode='train',
+ validation_in_fit=False,
**kwargs):
"""Loop function for arrays of data with modes 'train'/'test'/'predict'.
@@ -164,6 +165,9 @@
validation_steps: Number of steps to run validation for (only if doing
validation from data tensors). Ignored with the default value of `None`.
mode: One of 'train'/'test'/'predict'.
+ validation_in_fit: if true, then this method is invoked from within
+ training iteration (for validation). In this case, do not copy weights
+ when using a tf.distribute.Strategy.
**kwargs: Additional arguments for backwards compatibility.
Returns:
@@ -230,21 +234,21 @@
aggregator = training_utils.MetricsAggregator(use_steps,
num_samples_or_steps)
- if model._distribution_strategy:
- training_distributed._copy_weights_to_distributed_model(model)
+ if model._distribution_strategy and not validation_in_fit:
+ training_distributed._copy_weights_to_distributed_model(
+ model, model._grouped_model)
callbacks.model.stop_training = False
callbacks._call_begin_hook(mode)
progbar.on_train_begin()
+
for epoch in range(initial_epoch, epochs):
if callbacks.model.stop_training:
break
# Setup work for each epoch
epoch_logs = {}
- if hasattr(model, 'metrics'):
- for m in model.metrics:
- m.reset_states()
+ model.reset_metrics()
callbacks.on_epoch_begin(epoch, epoch_logs, mode=mode)
progbar.on_epoch_begin(epoch, epoch_logs)
@@ -356,7 +360,8 @@
steps_per_epoch=validation_steps,
callbacks=callbacks,
verbose=0,
- mode='test')
+ mode='test',
+ validation_in_fit=True)
if not isinstance(val_results, list):
val_results = [val_results]
epoch_logs.update(
@@ -367,7 +372,11 @@
callbacks._call_end_hook(mode)
if model._distribution_strategy:
- training_distributed._copy_weights_to_original_model(model, mode)
+ # TODO(priyag, psv): Copy back metrics to the original model as well?
+ if not validation_in_fit:
+ training_distributed._copy_weights_to_original_model(
+ model, model._grouped_model, mode)
+
scope.__exit__(None, None, None)
if mode == 'train':
diff --git a/tensorflow/python/keras/engine/training_dataset_test.py b/tensorflow/python/keras/engine/training_dataset_test.py
index 13e4abf..d6cc93d 100644
--- a/tensorflow/python/keras/engine/training_dataset_test.py
+++ b/tensorflow/python/keras/engine/training_dataset_test.py
@@ -24,6 +24,7 @@
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util as tf_test_util
from tensorflow.python.keras import keras_parameterized
@@ -37,7 +38,8 @@
class TestTrainingWithDatasetIterators(keras_parameterized.TestCase):
- @keras_parameterized.run_all_keras_modes_with_all_model_types
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_training_and_eval_methods_on_iterators_single_io(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -97,7 +99,8 @@
'you should specify the `steps` argument'):
model.predict(iterator, verbose=0)
- @keras_parameterized.run_all_keras_modes_with_all_model_types
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_get_next_op_created_once(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -119,7 +122,8 @@
ops.get_default_graph().finalize()
model.fit(iterator, epochs=1, steps_per_epoch=2, verbose=1)
- @keras_parameterized.run_all_keras_modes_with_all_model_types
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_iterators_running_out_of_data(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -144,12 +148,14 @@
class TestTrainingWithDataset(keras_parameterized.TestCase):
- # TODO(kaftan) Run w/ all model types.
- # Seems like subclass models has a bug, file ticket
- @keras_parameterized.run_all_keras_modes_with_all_model_types(
- exclude_models='subclass'
- )
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_calling_model_on_same_dataset(self):
+ if ((not testing_utils.should_run_eagerly())
+ and testing_utils.get_model_type() == 'subclass'
+ and context.executing_eagerly()):
+ self.skipTest('b/120673224')
+
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
loss = 'mse'
@@ -172,7 +178,8 @@
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0,
validation_data=dataset, validation_steps=2)
- @keras_parameterized.run_all_keras_modes_with_all_model_types
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_training_and_eval_methods_on_dataset(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -231,11 +238,8 @@
'you should specify the `steps` argument'):
model.predict(dataset, verbose=0)
- # TODO(kaftan) Run w/ all model types.
- # Seems like subclass models has a bug, file ticket
- @keras_parameterized.run_all_keras_modes_with_all_model_types(
- exclude_models='subclass'
- )
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_dataset_with_sample_weights(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -256,7 +260,8 @@
model.evaluate(dataset, steps=2, verbose=1)
model.predict(dataset, steps=2)
- @keras_parameterized.run_all_keras_modes_with_all_model_types
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_dataset_with_sparse_labels(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -305,11 +310,8 @@
class TestMetricsWithDatasetIterators(keras_parameterized.TestCase):
- # TODO(kaftan) Run w/ all model types.
- # Seems like subclass models has a bug, file ticket
- @keras_parameterized.run_all_keras_modes_with_all_model_types(
- exclude_models='subclass'
- )
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def test_metrics_correctness_with_iterator(self):
layers = [
keras.layers.Dense(8, activation='relu', input_dim=4,
diff --git a/tensorflow/python/keras/engine/training_distributed.py b/tensorflow/python/keras/engine/training_distributed.py
index 473f06d..ffb0266 100644
--- a/tensorflow/python/keras/engine/training_distributed.py
+++ b/tensorflow/python/keras/engine/training_distributed.py
@@ -163,11 +163,9 @@
do_validation = bool(validation_steps)
# Copy the weights from the original model to each of the replicated models.
- orig_model_weights = model.get_weights()
with current_strategy.scope():
- distributed_model = current_strategy.unwrap(model._grouped_model_train)[0]
- distributed_training_utils.set_weights(
- current_strategy, distributed_model, orig_model_weights)
+ _copy_weights_to_distributed_model(model, model._grouped_model_train)
+
callbacks = cbks.configure_callbacks(
callbacks,
model,
@@ -185,6 +183,8 @@
callbacks.on_train_begin()
for epoch in range(initial_epoch, epochs):
+ with current_strategy.scope():
+ _reset_metrics(model, model._grouped_model_train)
callbacks.on_epoch_begin(epoch)
epoch_logs = {}
step_index = 0
@@ -217,9 +217,8 @@
# Since we create a new clone from the original model we need to copy
# the weights back to the original model before we can run validation.
with current_strategy.scope():
- updated_weights = current_strategy.unwrap(
- model._grouped_model_train)[0].get_weights()
- model.set_weights(updated_weights)
+ _copy_weights_to_original_model(model, model._grouped_model_train,
+ 'train')
val_outs = experimental_test_loop( # pylint: disable=undefined-variable
model,
@@ -240,9 +239,7 @@
# Copy the weights back from the replicated model to the original model.
with current_strategy.scope():
- updated_weights = current_strategy.unwrap(
- model._grouped_model_train)[0].get_weights()
- model.set_weights(updated_weights)
+ _copy_weights_to_original_model(model, model._grouped_model_train, 'train')
K.get_session().run(current_strategy.finalize())
return model.history
@@ -345,22 +342,26 @@
progbar = Progbar(target=steps)
# Copy the weights from the original model to each of the replicated models.
- orig_model_weights = model.get_weights()
with current_strategy.scope():
- distributed_model = current_strategy.unwrap(model._grouped_model_test)[0]
- distributed_training_utils.set_weights(
- current_strategy, distributed_model, orig_model_weights)
-
+ _copy_weights_to_distributed_model(model, model._grouped_model_test)
+ _reset_metrics(model, model._grouped_model_test)
assert steps is not None
outs = [0.] * len(model.metrics_names)
for step in range(steps):
_, batch_outs = K.get_session().run([test_op, output_tensors])
for i, label in enumerate(model.metrics_names):
- outs[i] += batch_outs[label]
+ if i == 0:
+ # Loss is stateless metrics.
+ outs[i] += batch_outs[label]
+ else:
+ # For all stateful metrics, the aggregation is handled by mirrored vars.
+ outs[i] = batch_outs[label]
+
if verbose >= 1:
progbar.update(step + 1)
- for i in range(len(outs)):
- outs[i] /= (steps)
+
+ if len(outs) >= 0:
+ outs[0] /= (steps)
if initialize_finalize_strategy:
K.get_session().run(current_strategy.finalize())
@@ -455,12 +456,9 @@
progbar = Progbar(target=steps)
# Copy the weights from the original model to each of the replicated models.
- orig_model_weights = model.get_weights()
with current_strategy.scope():
- distributed_model = current_strategy.unwrap(model._grouped_model_predict)[0]
- distributed_training_utils.set_weights(
- current_strategy, distributed_model, orig_model_weights)
-
+ _copy_weights_to_distributed_model(model, model._grouped_model_predict)
+ _reset_metrics(model, model._grouped_model_predict)
assert steps is not None
# Since we do not know how many samples we will see, we cannot pre-allocate
# the returned Numpy arrays. Instead, we store one array per batch seen
@@ -572,11 +570,8 @@
# Validate that all the elements in x and y are of the same type and shape.
# We can then pass the first element of x and y to `_standardize_weights`
# below and be confident of the output.
- x_values, y_values, sample_weights_values = distributed_training_utils.\
- validate_distributed_dataset_inputs(model._distribution_strategy, x, y,
- sample_weights)
- model._standardize_weights(x_values, y_values,
- sample_weight=sample_weights_values)
+ distributed_training_utils.validate_distributed_dataset_inputs(
+ model._distribution_strategy, x, y, sample_weights)
return x, y, sample_weights
@@ -695,22 +690,23 @@
return ins
-def _copy_weights_to_distributed_model(model):
+def _copy_weights_to_distributed_model(original_model, grouped_model):
"""Copies weights from original model to distributed models."""
- if model._distribution_strategy:
- # Copy the weights from the original model to each of the replicated models.
- orig_model_weights = model.get_weights()
- distributed_model = model._distribution_strategy.unwrap(
- model._grouped_model)[0]
- distributed_training_utils.set_weights(
- model._distribution_strategy, distributed_model, orig_model_weights)
+ strategy = original_model._distribution_strategy
+ if strategy:
+ # Copy the weights from the original model to each of the replicated
+ # models.
+ orig_model_weights = original_model.get_weights()
+ distributed_model = strategy.unwrap(grouped_model)[0]
+ distributed_training_utils.set_weights(strategy, distributed_model,
+ orig_model_weights)
-def _copy_weights_to_original_model(model, mode):
+def _copy_weights_to_original_model(model, grouped_model, mode):
"""Copies weights from first distributed model back to original model."""
if model._distribution_strategy and mode == 'train':
updated_weights = model._distribution_strategy.unwrap(
- model._grouped_model)[0].get_weights()
+ grouped_model)[0].get_weights()
model.set_weights(updated_weights)
@@ -724,3 +720,11 @@
total_batch_outs.append(np.concatenate(nest.flatten(nested_outs)))
return total_batch_outs
return batch_outs
+
+
+def _reset_metrics(model, distributed_model=None):
+ if model._distribution_strategy:
+ distributed_model = (
+ distributed_model or
+ model._distribution_strategy.unwrap(model._grouped_model)[0])
+ distributed_model.reset_metrics()
diff --git a/tensorflow/python/keras/engine/training_eager_test.py b/tensorflow/python/keras/engine/training_eager_test.py
index 3fabbb1..f95a502 100644
--- a/tensorflow/python/keras/engine/training_eager_test.py
+++ b/tensorflow/python/keras/engine/training_eager_test.py
@@ -246,6 +246,21 @@
layer(1.) # Plain-value inputs are only valid in eager mode.
self.assertEqual(1, len(layer.losses))
+ def test_predict_correctness(self):
+ i1 = keras.layers.Input(shape=(4, 5))
+ i2 = keras.layers.Input(shape=(4, 5))
+ i3 = keras.layers.Input(shape=(4, 5))
+ o = keras.layers.add([i1, i2, i3])
+ model = keras.models.Model([i1, i2, i3], o)
+ model.run_eagerly = True
+
+ x1 = np.random.random((2, 4, 5))
+ x2 = np.random.random((2, 4, 5))
+ x3 = np.random.random((2, 4, 5))
+ out = model.predict([x1, x2, x3])
+
+ self.assertAllClose(out, x1 + x2 + x3)
+
if __name__ == '__main__':
ops.enable_eager_execution()
diff --git a/tensorflow/python/keras/engine/training_generator.py b/tensorflow/python/keras/engine/training_generator.py
index 0abf0b8..88f21b3 100644
--- a/tensorflow/python/keras/engine/training_generator.py
+++ b/tensorflow/python/keras/engine/training_generator.py
@@ -49,7 +49,7 @@
max_queue_size=10,
workers=1,
use_multiprocessing=False,
- shuffle=True,
+ shuffle=False,
initial_epoch=0,
mode='train',
batch_size=None,
@@ -246,8 +246,10 @@
# Maintain compatibility with the existing names.
fit_generator = functools.partial(model_iteration, mode='train')
-evaluate_generator = functools.partial(model_iteration, mode='test')
-predict_generator = functools.partial(model_iteration, mode='predict')
+evaluate_generator = functools.partial(
+ model_iteration, mode='test', shuffle=False)
+predict_generator = functools.partial(
+ model_iteration, mode='predict', shuffle=False)
def _get_next_batch(output_generator, mode):
diff --git a/tensorflow/python/keras/engine/training_generator_test.py b/tensorflow/python/keras/engine/training_generator_test.py
index 8941428..956ca2f 100644
--- a/tensorflow/python/keras/engine/training_generator_test.py
+++ b/tensorflow/python/keras/engine/training_generator_test.py
@@ -29,6 +29,7 @@
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.eager import context
from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import metrics as metrics_module
from tensorflow.python.keras import testing_utils
from tensorflow.python.keras.engine import training_generator
@@ -60,20 +61,17 @@
yield x, y, w
-@tf_test_util.run_all_in_graph_and_eager_modes
-class TestGeneratorMethods(test.TestCase, parameterized.TestCase):
+class TestGeneratorMethods(keras_parameterized.TestCase):
@unittest.skipIf(
os.name == 'nt',
'use_multiprocessing=True does not work on windows properly.')
- @parameterized.parameters('sequential', 'functional')
- def test_fit_generator_method(self, model_type):
- if model_type == 'sequential':
- model = testing_utils.get_small_sequential_mlp(
- num_hidden=3, num_classes=4, input_dim=2)
- else:
- model = testing_utils.get_small_functional_mlp(
- num_hidden=3, num_classes=4, input_dim=2)
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
+ def test_fit_generator_method(self):
+ model = testing_utils.get_small_mlp(
+ num_hidden=3, num_classes=4, input_dim=2)
model.compile(
loss='mse',
optimizer='sgd',
@@ -109,19 +107,17 @@
@unittest.skipIf(
os.name == 'nt',
'use_multiprocessing=True does not work on windows properly.')
- @parameterized.parameters('sequential', 'functional')
- def test_evaluate_generator_method(self, model_type):
- if model_type == 'sequential':
- model = testing_utils.get_small_sequential_mlp(
- num_hidden=3, num_classes=4, input_dim=2)
- else:
- model = testing_utils.get_small_functional_mlp(
- num_hidden=3, num_classes=4, input_dim=2)
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
+ def test_evaluate_generator_method(self):
+ model = testing_utils.get_small_mlp(
+ num_hidden=3, num_classes=4, input_dim=2)
model.compile(
loss='mse',
optimizer='sgd',
- metrics=['mae', metrics_module.CategoricalAccuracy()])
- model.summary()
+ metrics=['mae', metrics_module.CategoricalAccuracy()],
+ run_eagerly=testing_utils.should_run_eagerly())
model.evaluate_generator(custom_generator(),
steps=5,
@@ -142,18 +138,16 @@
@unittest.skipIf(
os.name == 'nt',
'use_multiprocessing=True does not work on windows properly.')
- @parameterized.parameters('sequential', 'functional')
- def test_predict_generator_method(self, model_type):
- if model_type == 'sequential':
- model = testing_utils.get_small_sequential_mlp(
- num_hidden=3, num_classes=4, input_dim=2)
- else:
- model = testing_utils.get_small_functional_mlp(
- num_hidden=3, num_classes=4, input_dim=2)
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
+ def test_predict_generator_method(self):
+ model = testing_utils.get_small_mlp(
+ num_hidden=3, num_classes=4, input_dim=2)
model.compile(
loss='mse',
optimizer='sgd',
- metrics=['mae', metrics_module.CategoricalAccuracy()])
+ metrics=['mae', metrics_module.CategoricalAccuracy()],
+ run_eagerly=testing_utils.should_run_eagerly())
model.predict_generator(custom_generator(),
steps=5,
@@ -183,13 +177,17 @@
max_queue_size=10,
workers=0)
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
def test_generator_methods_with_sample_weights(self):
- model = keras.models.Sequential()
- model.add(keras.layers.Dense(4, input_shape=(2,)))
+ model = testing_utils.get_small_mlp(
+ num_hidden=3, num_classes=4, input_dim=2)
model.compile(
loss='mse',
optimizer='sgd',
- metrics=['mae', metrics_module.CategoricalAccuracy()])
+ metrics=['mae', metrics_module.CategoricalAccuracy()],
+ run_eagerly=testing_utils.should_run_eagerly())
model.fit_generator(custom_generator(mode=3),
steps_per_epoch=5,
@@ -214,15 +212,19 @@
max_queue_size=10,
use_multiprocessing=False)
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
def test_generator_methods_invalid_use_case(self):
def invalid_generator():
while 1:
yield 0
- model = keras.models.Sequential()
- model.add(keras.layers.Dense(4, input_shape=(2,)))
- model.compile(loss='mse', optimizer='sgd')
+ model = testing_utils.get_small_mlp(
+ num_hidden=3, num_classes=4, input_dim=2)
+ model.compile(loss='mse', optimizer='sgd',
+ run_eagerly=testing_utils.should_run_eagerly())
with self.assertRaises(ValueError):
model.fit_generator(invalid_generator(),
@@ -251,6 +253,9 @@
max_queue_size=10,
use_multiprocessing=False)
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
def test_generator_input_to_fit_eval_predict(self):
val_data = np.ones([10, 10], np.float32), np.ones([10, 1], np.float32)
@@ -258,12 +263,11 @@
while True:
yield np.ones([10, 10], np.float32), np.ones([10, 1], np.float32)
- inputs = keras.layers.Input(shape=(10,))
- x = keras.layers.Dense(10, activation='relu')(inputs)
- outputs = keras.layers.Dense(1, activation='sigmoid')(x)
- model = keras.Model(inputs, outputs)
+ model = testing_utils.get_small_mlp(
+ num_hidden=10, num_classes=1, input_dim=10)
- model.compile(RMSPropOptimizer(0.001), 'binary_crossentropy')
+ model.compile(RMSPropOptimizer(0.001), 'binary_crossentropy',
+ run_eagerly=testing_utils.should_run_eagerly())
model.fit(
ones_generator(),
steps_per_epoch=2,
@@ -273,9 +277,11 @@
model.predict(ones_generator(), steps=2)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class TestGeneratorMethodsWithSequences(test.TestCase):
+class TestGeneratorMethodsWithSequences(keras_parameterized.TestCase):
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
def test_training_with_sequences(self):
class DummySequence(keras.utils.Sequence):
@@ -286,8 +292,8 @@
def __len__(self):
return 10
- model = keras.models.Sequential()
- model.add(keras.layers.Dense(4, input_shape=(2,)))
+ model = testing_utils.get_small_mlp(
+ num_hidden=3, num_classes=4, input_dim=2)
model.compile(loss='mse', optimizer='sgd')
model.fit_generator(DummySequence(),
@@ -305,6 +311,9 @@
workers=0,
use_multiprocessing=False)
+ # TODO(b/120940700): Bug with subclassed model inputs.
+ @keras_parameterized.run_with_all_model_types(exclude_models='subclass')
+ @keras_parameterized.run_all_keras_modes
def test_sequence_input_to_fit_eval_predict(self):
val_data = np.ones([10, 10], np.float32), np.ones([10, 1], np.float32)
@@ -316,10 +325,8 @@
def __len__(self):
return 2
- inputs = keras.layers.Input(shape=(10,))
- x = keras.layers.Dense(10, activation='relu')(inputs)
- outputs = keras.layers.Dense(1, activation='sigmoid')(x)
- model = keras.Model(inputs, outputs)
+ model = testing_utils.get_small_mlp(
+ num_hidden=10, num_classes=1, input_dim=10)
model.compile(RMSPropOptimizer(0.001), 'binary_crossentropy')
model.fit(CustomSequence(), validation_data=val_data, epochs=2)
diff --git a/tensorflow/python/keras/engine/training_test.py b/tensorflow/python/keras/engine/training_test.py
index 572cde1..9d56eb2 100644
--- a/tensorflow/python/keras/engine/training_test.py
+++ b/tensorflow/python/keras/engine/training_test.py
@@ -53,9 +53,8 @@
class TrainingTest(keras_parameterized.TestCase):
- @keras_parameterized.run_all_keras_modes_with_all_model_types(
- exclude_models='sequential'
- )
+ @keras_parameterized.run_with_all_model_types(exclude_models='sequential')
+ @keras_parameterized.run_all_keras_modes
def test_fit_on_arrays(self):
input_a = keras.layers.Input(shape=(3,), name='input_a')
input_b = keras.layers.Input(shape=(3,), name='input_b')
@@ -251,8 +250,10 @@
run_eagerly=testing_utils.should_run_eagerly())
# This will work
model.fit([input_a_np], output_d_np, epochs=1)
- with self.assertRaises(ValueError):
- model.fit([input_a_np, input_a_np], output_d_np, epochs=1)
+ # TODO(gsundeep) Test only works in eager, file ticket
+ if testing_utils.should_run_eagerly() and context.executing_eagerly():
+ with self.assertRaises(ValueError):
+ model.fit([input_a_np, input_a_np], output_d_np, epochs=1)
# Test model on a list of floats
input_a_np = np.random.random((10, 3))
@@ -472,7 +473,7 @@
metrics=['accuracy'],
run_eagerly=testing_utils.should_run_eagerly())
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_that_trainable_disables_updates(self):
val_a = np.random.random((10, 4))
val_out = np.random.random((10, 4))
@@ -797,8 +798,9 @@
def test_class_weights(self):
num_classes = 5
batch_size = 5
- epochs = 5
+ epochs = 10
weighted_class = 3
+ weight = 10.
train_samples = 1000
test_samples = 1000
input_dim = 5
@@ -827,7 +829,7 @@
test_ids = np.where(int_y_test == np.array(weighted_class))[0]
class_weight = dict([(i, 1.) for i in range(num_classes)])
- class_weight[weighted_class] = 2.
+ class_weight[weighted_class] = weight
sample_weight = np.ones((y_train.shape[0]))
sample_weight[int_y_train == weighted_class] = 2.
@@ -867,8 +869,9 @@
def test_sample_weights(self):
num_classes = 5
batch_size = 5
- epochs = 5
+ epochs = 10
weighted_class = 3
+ weight = 10.
train_samples = 1000
test_samples = 1000
input_dim = 5
@@ -897,7 +900,7 @@
test_ids = np.where(int_y_test == np.array(weighted_class))[0]
sample_weight = np.ones((y_train.shape[0]))
- sample_weight[int_y_train == weighted_class] = 2.
+ sample_weight[int_y_train == weighted_class] = weight
model.fit(
x_train,
@@ -964,8 +967,9 @@
def test_temporal_sample_weights(self):
num_classes = 5
batch_size = 5
- epochs = 5
+ epochs = 10
weighted_class = 3
+ weight = 10.
train_samples = 1000
test_samples = 1000
input_dim = 5
@@ -994,7 +998,7 @@
test_ids = np.where(int_y_test == np.array(weighted_class))[0]
sample_weight = np.ones((y_train.shape[0]))
- sample_weight[int_y_train == weighted_class] = 2.
+ sample_weight[int_y_train == weighted_class] = weight
temporal_x_train = np.reshape(x_train, (len(x_train), 1,
x_train.shape[1]))
@@ -1015,7 +1019,7 @@
model.compile(
RMSPropOptimizer(learning_rate=learning_rate),
- loss='binary_crossentropy',
+ loss='categorical_crossentropy',
metrics=['acc', metrics_module.CategoricalAccuracy()],
weighted_metrics=['mae', metrics_module.CategoricalAccuracy()],
sample_weight_mode='temporal',
@@ -1282,7 +1286,7 @@
class TestDynamicTrainability(keras_parameterized.TestCase):
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_trainable_warning(self):
with self.cached_session():
x = np.random.random((5, 3))
@@ -1296,7 +1300,7 @@
model.train_on_batch(x, y)
self.assertRaises(Warning)
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_trainable_argument(self):
with self.cached_session():
x = np.random.random((5, 3))
diff --git a/tensorflow/python/keras/engine/training_utils.py b/tensorflow/python/keras/engine/training_utils.py
index 9301302..ec6b397 100644
--- a/tensorflow/python/keras/engine/training_utils.py
+++ b/tensorflow/python/keras/engine/training_utils.py
@@ -27,9 +27,11 @@
from tensorflow.python.data.ops import iterator_ops
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 tensor_shape
+from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import tensor_util
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import callbacks as cbks
@@ -232,10 +234,14 @@
return None # Edge case where ins == [static_learning_phase]
-def standardize_single_array(x):
+def standardize_single_array(x, expected_shape=None):
+ """Expand data of shape (x,) to (x, 1), unless len(expected_shape)==1."""
if x is None:
return None
- if x.shape is not None and len(x.shape) == 1:
+
+ if (x.shape is not None
+ and len(x.shape) == 1
+ and (expected_shape is None or len(expected_shape) != 1)):
if tensor_util.is_tensor(x):
x = array_ops.expand_dims(x, axis=1)
else:
@@ -301,7 +307,11 @@
else:
data = data.values if data.__class__.__name__ == 'DataFrame' else data
data = [data]
- data = [standardize_single_array(x) for x in data]
+ if shapes is not None:
+ data = [standardize_single_array(x, shape)
+ for (x, shape) in zip(data, shapes)]
+ else:
+ data = [standardize_single_array(x) for x in data]
if len(data) != len(names):
if data and hasattr(data[0], 'shape'):
@@ -1183,3 +1193,61 @@
if batch_input_shape is not None:
return tensor_shape.as_dimension(batch_input_shape[0]).value
return None
+
+
+def generic_output_names(outputs_list):
+ return ['output_%d' % (i + 1) for i in range(len(outputs_list))]
+
+
+def trace_model_call(model, input_signature=None):
+ """Trace the model call to create a tf.function for exporting a Keras model.
+
+ Args:
+ model: A Keras model.
+ input_signature: optional, a list of tf.TensorSpec objects specifying the
+ inputs to the model.
+
+ Returns:
+ A tf.function wrapping the model's call function with input signatures set.
+
+ Raises:
+ ValueError: if input signature cannot be inferred from the model.
+ """
+ if input_signature is None:
+ if isinstance(model.call, def_function.PolymorphicFunction):
+ input_signature = model.call.input_signature
+
+ if input_signature is None:
+ try:
+ inputs = model.inputs
+ input_names = model.input_names
+ except AttributeError:
+ raise ValueError(
+ 'Model {} cannot be saved because the input shapes have not been '
+ 'set. Usually, input shapes are automatically determined from calling'
+ ' .fit() or .predict(). To manually set the shapes, call '
+ 'model._set_inputs(inputs).'.format(model))
+ input_specs = []
+ for input_tensor, input_name in zip(inputs, input_names):
+ input_specs.append(tensor_spec.TensorSpec(
+ shape=input_tensor.shape, dtype=input_tensor.dtype,
+ name=input_name))
+ # The input signature of the call function is a list with one element, since
+ # all tensor inputs must be passed in as the first argument.
+ input_signature = [input_specs] if len(input_specs) > 1 else input_specs
+
+ @def_function.function(input_signature=input_signature)
+ def _wrapped_model(*args):
+ """A concrete tf.function that wraps the model's call function."""
+ # When given a single input, Keras models will call the model on the tensor
+ # rather than a list consisting of the single tensor.
+ inputs = args[0] if len(input_signature) == 1 else list(args)
+ outputs_list = nest.flatten(model(inputs=inputs))
+ try:
+ output_names = model.output_names
+ except AttributeError:
+ output_names = generic_output_names(outputs_list)
+ return {name: output for name, output in zip(output_names, outputs_list)}
+
+ return _wrapped_model
+
diff --git a/tensorflow/python/keras/engine/training_utils_test.py b/tensorflow/python/keras/engine/training_utils_test.py
index 44ea239..d8acec3 100644
--- a/tensorflow/python/keras/engine/training_utils_test.py
+++ b/tensorflow/python/keras/engine/training_utils_test.py
@@ -18,13 +18,30 @@
from __future__ import division
from __future__ import print_function
+import os
+
import numpy as np
+
+from tensorflow.python.client import session as session_lib
+from tensorflow.python import keras
from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+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.keras import backend as K
+from tensorflow.python.keras import keras_parameterized
+from tensorflow.python.keras import testing_utils
from tensorflow.python.keras.engine import training_utils
from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
+from tensorflow.python.saved_model import loader
+from tensorflow.python.saved_model import save as save_lib
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import tag_constants
class ModelInputsTest(test.TestCase):
@@ -85,5 +102,168 @@
self.assertTrue(tf_utils.is_symbolic_tensor(vals['b']))
+class TraceModelCallTest(keras_parameterized.TestCase):
+
+ def _assert_all_close(self, expected, actual):
+ if not context.executing_eagerly():
+ with self.cached_session() as sess:
+ K._initialize_variables(sess)
+ self.assertAllClose(expected, actual)
+ else:
+ self.assertAllClose(expected, actual)
+
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
+ def test_trace_model_outputs(self):
+ input_dim = 5 if testing_utils.get_model_type() == 'functional' else None
+ model = testing_utils.get_small_mlp(10, 3, input_dim)
+ inputs = array_ops.ones((8, 5))
+
+ if input_dim is None:
+ with self.assertRaisesRegexp(ValueError,
+ 'input shapes have not been set'):
+ training_utils.trace_model_call(model)
+ model._set_inputs(inputs)
+
+ fn = training_utils.trace_model_call(model)
+ signature_outputs = fn(inputs)
+ expected_outputs = {model.output_names[0]: model(inputs)}
+
+ self._assert_all_close(expected_outputs, signature_outputs)
+
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
+ def test_trace_model_outputs_after_fitting(self):
+ input_dim = 5 if testing_utils.get_model_type() == 'functional' else None
+ model = testing_utils.get_small_mlp(10, 3, input_dim)
+ model.compile(optimizer='sgd', loss='mse')
+ model.fit(x=np.random.random((8, 5)),
+ y=np.random.random((8, 3)), epochs=2)
+
+ inputs = array_ops.ones((8, 5))
+
+ fn = training_utils.trace_model_call(model)
+ signature_outputs = fn(inputs)
+ expected_outputs = {model.output_names[0]: model(inputs)}
+
+ self._assert_all_close(expected_outputs, signature_outputs)
+
+ @keras_parameterized.run_with_all_model_types(exclude_models='sequential')
+ @keras_parameterized.run_all_keras_modes
+ def test_trace_multi_io_model_outputs(self):
+ input_dim = 5
+ num_classes = 3
+ num_classes_b = 4
+ input_a = keras.layers.Input(shape=(input_dim,), name='input_a')
+ input_b = keras.layers.Input(shape=(input_dim,), name='input_b')
+
+ dense = keras.layers.Dense(num_classes, name='dense')
+ dense2 = keras.layers.Dense(num_classes_b, name='dense2')
+ dropout = keras.layers.Dropout(0.5, name='dropout')
+ branch_a = [input_a, dense]
+ branch_b = [input_b, dense, dense2, dropout]
+
+ model = testing_utils.get_multi_io_model(branch_a, branch_b)
+
+ input_a_np = np.random.random((10, input_dim)).astype(np.float32)
+ input_b_np = np.random.random((10, input_dim)).astype(np.float32)
+
+ if testing_utils.get_model_type() == 'subclass':
+ with self.assertRaisesRegexp(ValueError,
+ 'input shapes have not been set'):
+ training_utils.trace_model_call(model)
+
+ model.compile(optimizer='sgd', loss='mse')
+ model.fit(x=[np.random.random((8, input_dim)).astype(np.float32),
+ np.random.random((8, input_dim)).astype(np.float32)],
+ y=[np.random.random((8, num_classes)).astype(np.float32),
+ np.random.random((8, num_classes_b)).astype(np.float32)],
+ epochs=2)
+
+ fn = training_utils.trace_model_call(model)
+ signature_outputs = fn([input_a_np, input_b_np])
+ outputs = model([input_a_np, input_b_np])
+ expected_outputs = {model.output_names[0]: outputs[0],
+ model.output_names[1]: outputs[1]}
+
+ self._assert_all_close(expected_outputs, signature_outputs)
+
+ @keras_parameterized.run_all_keras_modes
+ def test_specify_input_signature(self):
+ model = testing_utils.get_small_sequential_mlp(10, 3, None)
+ inputs = array_ops.ones((8, 5))
+
+ with self.assertRaisesRegexp(ValueError, 'input shapes have not been set'):
+ training_utils.trace_model_call(model)
+
+ fn = training_utils.trace_model_call(
+ model, [tensor_spec.TensorSpec(shape=[None, 5], dtype=dtypes.float32)])
+ signature_outputs = fn(inputs)
+ expected_outputs = {model.output_names[0]: model(inputs)}
+ self._assert_all_close(expected_outputs, signature_outputs)
+
+ @keras_parameterized.run_all_keras_modes
+ def test_subclassed_model_with_input_signature(self):
+
+ class Model(keras.Model):
+
+ def __init__(self):
+ super(Model, self).__init__()
+ self.dense = keras.layers.Dense(3, name='dense')
+
+ @def_function.function(
+ input_signature=[[tensor_spec.TensorSpec([None, 5], dtypes.float32),
+ tensor_spec.TensorSpec([None], dtypes.float32)]],)
+ def call(self, inputs, *args):
+ x, y = inputs
+ return self.dense(x) + y
+
+ model = Model()
+ fn = training_utils.trace_model_call(model)
+ x = array_ops.ones((8, 5), dtype=dtypes.float32)
+ y = array_ops.ones((3,), dtype=dtypes.float32)
+ expected_outputs = {'output_1': model([x, y])}
+ signature_outputs = fn([x, y])
+ self._assert_all_close(expected_outputs, signature_outputs)
+
+
+def _import_and_infer(save_dir, inputs):
+ """Import a SavedModel into a TF 1.x-style graph and run `signature_key`."""
+ graph = ops.Graph()
+ with graph.as_default(), session_lib.Session() as session:
+ model = loader.load(session, [tag_constants.SERVING], save_dir)
+ signature = model.signature_def[
+ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+ assert set(inputs.keys()) == set(signature.inputs.keys())
+ feed_dict = {}
+ for arg_name in inputs.keys():
+ feed_dict[graph.get_tensor_by_name(signature.inputs[arg_name].name)] = (
+ inputs[arg_name])
+ output_dict = {}
+ for output_name, output_tensor_info in signature.outputs.items():
+ output_dict[output_name] = graph.get_tensor_by_name(
+ output_tensor_info.name)
+ return session.run(output_dict, feed_dict=feed_dict)
+
+
+class ModelSaveTest(keras_parameterized.TestCase):
+
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
+ def test_model_save(self):
+ input_dim = 5
+ model = testing_utils.get_small_mlp(10, 3, input_dim)
+ inputs = array_ops.ones((8, 5))
+
+ if testing_utils.get_model_type() == 'subclass':
+ model._set_inputs(inputs)
+
+ save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
+ save_lib.save(model, save_dir)
+
+ self.assertAllClose(
+ {model.output_names[0]: model.predict_on_batch(inputs)},
+ _import_and_infer(save_dir, {model.input_names[0]: np.ones((8, 5))}))
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/integration_test.py b/tensorflow/python/keras/integration_test.py
index f1a0932..8d65f63 100644
--- a/tensorflow/python/keras/integration_test.py
+++ b/tensorflow/python/keras/integration_test.py
@@ -35,7 +35,7 @@
def test_version(self):
self.assertTrue(keras.__version__.endswith('-tf'))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_vector_classification_sequential(self):
with self.cached_session():
np.random.seed(1337)
@@ -168,7 +168,7 @@
verbose=2)
self.assertGreater(history.history['val_acc'][-1], 0.7)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_video_classification_functional(self):
with self.cached_session():
np.random.seed(1337)
@@ -197,7 +197,7 @@
verbose=2)
self.assertGreater(history.history['val_acc'][-1], 0.7)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_vector_classification_shared_sequential(self):
# Test that Sequential models that feature internal updates
# and internal losses can be shared.
@@ -232,7 +232,7 @@
verbose=2)
self.assertGreater(history.history['val_acc'][-1], 0.7)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_vector_classification_shared_model(self):
# Test that functional models that feature internal updates
# and internal losses can be shared.
diff --git a/tensorflow/python/keras/keras_parameterized.py b/tensorflow/python/keras/keras_parameterized.py
index 57e66d8..d76bbad 100644
--- a/tensorflow/python/keras/keras_parameterized.py
+++ b/tensorflow/python/keras/keras_parameterized.py
@@ -18,6 +18,11 @@
from __future__ import division
from __future__ import print_function
+import collections
+import functools
+import itertools
+import unittest
+
from absl.testing import parameterized
from tensorflow.python import keras
@@ -38,21 +43,20 @@
# TODO(kaftan): Possibly enable 'subclass_custom_build' when tests begin to pass
# it. Or perhaps make 'subclass' always use a custom build method.
def run_with_all_model_types(
- func=None,
- exclude_models=None,
- extra_parameters=None):
+ test_or_class=None,
+ exclude_models=None):
"""Execute the decorated test with all Keras model types.
- This decorator is intended to be applied to test methods in
- a `testing_utils.KerasTestCase` class. Doing so will cause the contents of
- the test method to be executed multiple times - once for each Keras model
- type.
+ This decorator is intended to be applied either to individual test methods in
+ a `keras_parameterized.TestCase` class, or directly to a test class that
+ extends it. Doing so will cause the contents of the individual test
+ method (or all test methods in the class) to be executed multiple times - once
+ for each Keras model type.
The Keras model types are: ['functional', 'subclass', 'sequential']
- Note: This decorator does not stack with any other parameterizing
- decorators. Use the extra_parameters arg if you would like to use additional
- parameterization.
+ Note: if stacking this decorator with absl.testing's parameterized decorators,
+ those should be at the bottom of the stack.
Various methods in `testing_utils` to get models will auto-generate a model
of the currently active Keras model type. This allows unittests to confirm
@@ -64,7 +68,7 @@
class MyTests(testing_utils.KerasTestCase):
@testing_utils.run_with_all_model_types(
- model_types = ['functional', 'subclass'])
+ exclude_models = ['sequential'])
def test_foo(self):
model = testing_utils.get_small_mlp(1, 4, input_dim=3)
optimizer = RMSPropOptimizer(learning_rate=0.001)
@@ -87,19 +91,40 @@
This test tries building a small mlp as both a functional model and as a
subclass model.
+ We can also annotate the whole class if we want this to apply to all tests in
+ the class:
+ ```python
+ @testing_utils.run_with_all_model_types(exclude_models = ['sequential'])
+ class MyTests(testing_utils.KerasTestCase):
+
+ def test_foo(self):
+ model = testing_utils.get_small_mlp(1, 4, input_dim=3)
+ optimizer = RMSPropOptimizer(learning_rate=0.001)
+ loss = 'mse'
+ metrics = ['mae']
+ model.compile(optimizer, loss, metrics=metrics)
+
+ inputs = np.zeros((10, 3))
+ targets = np.zeros((10, 4))
+ dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
+ dataset = dataset.repeat(100)
+ dataset = dataset.batch(10)
+
+ model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
+
+ if __name__ == "__main__":
+ tf.test.main()
+ ```
+
Args:
- func: function to be annotated. If `func` is None, this method returns a
- decorator the can be applied to a function. If `func` is not None this
- returns the decorator applied to `func`.
+ test_or_class: test method or class to be annotated. If None,
+ this method returns a decorator that can be applied to a test method or
+ test class. If it is not None this returns the decorator applied to the
+ test or class.
exclude_models: A collection of Keras model types to not run.
(May also be a single model type not wrapped in a collection).
Defaults to None.
- extra_parameters: A list of dicts containing additional parameters for
- parameterizing the decorated test method. This additional parameters must
- be arguments in the test method's signature.
- See the absl.testing.parameterized decorators for more info on how these
- work.
Returns:
Returns a decorator that will run the decorated test method multiple times:
@@ -109,30 +134,44 @@
ImportError: If abseil parameterized is not installed or not included as
a target dependency.
"""
- return run_all_keras_modes_with_all_model_types(
- func, skip_run_modes=True, exclude_models=exclude_models,
- extra_parameters=extra_parameters)
+ model_types = ['functional', 'subclass', 'sequential']
+ params = [('_%s' % model, model) for model in model_types
+ if model not in nest.flatten(exclude_models)]
+
+ def single_method_decorator(f):
+ """Decorator that constructs the test cases."""
+ # Use named_parameters so it can be individually run from the command line
+ @parameterized.named_parameters(*params)
+ @functools.wraps(f)
+ def decorated(self, model_type, *args, **kwargs):
+ """A run of a single test case w/ the specified model type."""
+ with testing_utils.model_type_scope(model_type):
+ f(self, *args, **kwargs)
+
+ return decorated
+
+ return _test_or_class_decorator(test_or_class, single_method_decorator)
def run_all_keras_modes(
- func=None,
+ test_or_class=None,
config=None,
- always_skip_v1=False,
- extra_parameters=None):
+ always_skip_v1=False):
"""Execute the decorated test with all keras execution modes.
- This decorator is intended to be applied to test methods in
- a `testing_utils.KerasTestCase` class. Doing so will cause the contents of
- the test method to be executed three times - once executing in legacy graph
- mode, once running eagerly and with `should_run_eagerly` returning True, and
- once running eagerly with `should_run_eagerly` returning False.
+ This decorator is intended to be applied either to individual test methods in
+ a `keras_parameterized.TestCase` class, or directly to a test class that
+ extends it. Doing so will cause the contents of the individual test
+ method (or all test methods in the class) to be executed multiple times -
+ once executing in legacy graph mode, once running eagerly and with
+ `should_run_eagerly` returning True, and once running eagerly with
+ `should_run_eagerly` returning False.
If Tensorflow v2 behavior is enabled, legacy graph mode will be skipped, and
the test will only run twice.
- Note: This decorator does not stack with any other parameterizing
- decorators. Use the extra_parameters arg if you would like to use additional
- parameterization.
+ Note: if stacking this decorator with absl.testing's parameterized decorators,
+ those should be at the bottom of the stack.
For example, consider the following unittest:
@@ -164,17 +203,14 @@
three Keras execution modes.
Args:
- func: function to be annotated. If `func` is None, this method returns a
- decorator the can be applied to a function. If `func` is not None this
- returns the decorator applied to `func`.
+ test_or_class: test method or class to be annotated. If None,
+ this method returns a decorator that can be applied to a test method or
+ test class. If it is not None this returns the decorator applied to the
+ test or class.
config: An optional config_pb2.ConfigProto to use to configure the
session when executing graphs.
always_skip_v1: If True, does not try running the legacy graph mode even
when Tensorflow v2 behavior is not enabled.
- extra_parameters: A list of dicts containing additional parameters for
- parameterizing the decorated test method. This additional parameters must
- be arguments in the test method's signature. See the
- absl.testing.parameterized decorators for more info on how these work.
Returns:
Returns a decorator that will run the decorated test method multiple times.
@@ -183,166 +219,80 @@
ImportError: If abseil parameterized is not installed or not included as
a target dependency.
"""
- return run_all_keras_modes_with_all_model_types(
- func=func,
- skip_model_types=True,
- config=config,
- always_skip_v1=always_skip_v1,
- extra_parameters=extra_parameters)
+ params = [('_v2_eager', 'v2_eager'),
+ ('_v2_function', 'v2_function')]
+ if not (always_skip_v1 or tf2.enabled()):
+ params.append(('_v1_graph', 'v1_graph'))
-
-def run_all_keras_modes_with_all_model_types(
- func=None,
- skip_run_modes=False,
- skip_model_types=False,
- config=None,
- always_skip_v1=False,
- exclude_models=None,
- extra_parameters=None):
- """Execute the decorated test with all keras run modes and model types.
-
- This decorator is intended to be applied to test methods in
- a `testing_utils.KerasTestCase` class. Doing so will cause the contents of
- the test method to be executed once for each combination of Keras run mode
- and model type.
-
- Note: This decorator does not stack with any other parameterizing
- decorators. Use the extra_parameters arg if you would like to use additional
- parameterization.
-
- For example, consider the following unittest:
-
- ```python
- class MyTests(testing_utils.KerasTestCase):
-
- @testing_utils.run_all_keras_modes_with_all_model_types
- def test_foo(self):
- model = testing_utils.get_small_functional_mlp(1, 4, input_dim=3)
- optimizer = RMSPropOptimizer(learning_rate=0.001)
- loss = 'mse'
- metrics = ['mae']
- model.compile(optimizer, loss, metrics=metrics,
- run_eagerly=testing_utils.should_run_eagerly())
-
- inputs = np.zeros((10, 3))
- targets = np.zeros((10, 4))
- dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
- dataset = dataset.repeat(100)
- dataset = dataset.batch(10)
-
- model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
-
- if __name__ == "__main__":
- tf.test.main()
- ```
-
- This test will try compiling & fitting the small functional mlp using all
- combinations of the Keras execution modes and the Keras model types.
-
- The three Keras execution modes are v1 graph mode, v2 functions, and v2 eager.
- If Tensorflow v2 behavior is enabled, v1 graph mode will be skipped.
- The Keras model types are: ['functional', 'subclass', 'sequential']
-
- Args:
- func: function to be annotated. If `func` is None, this method returns a
- decorator the can be applied to a function. If `func` is not None this
- returns the decorator applied to `func`.
- skip_run_modes: If True, does not try testing different Keras run modes.
- Defaults to False, which tests all combinations of run modes.
- skip_model_types: If True, does not try testing different Keras model types.
- Defaults to False, which tests all combinations of model types.
- config: An optional config_pb2.ConfigProto to use to configure the
- session when executing graphs.
- always_skip_v1: If True, does not try running the legacy graph mode even
- when Tensorflow v2 behavior is not enabled.
- exclude_models: A collection of Keras model types to not run.
- (May also be a single model type not wrapped in a collection).
- extra_parameters: A list of dicts containing additional parameters for
- parameterizing the decorated test method. This additional parameters must
- be arguments in the test method's signature. See the
- absl.testing.parameterized decorators for more info on how these work.
-
- Returns:
- Returns a decorator that will run the decorated test method multiple times.
-
- Raises:
- ImportError: If abseil parameterized is not installed or not included as
- a target dependency.
- """
- if not parameterized:
- raise ImportError('To use the Keras parameterized testing utilities, you '
- 'must have absl.testing.parameterized installed and '
- 'included as a target dependency.')
- params = []
- if not skip_run_modes:
- if not (always_skip_v1 or tf2.enabled()):
- params.append({'testcase_name': '_v1_graph', '_run_mode': 'v1_graph'})
- params.append({'testcase_name': '_v2_eager', '_run_mode': 'v2_eager'})
- params.append({'testcase_name': '_v2_function', '_run_mode': 'v2_function'})
- else:
- params.append({'testcase_name': ''})
-
- if not skip_model_types:
- model_types = ['functional', 'subclass', 'sequential']
- old_params = params
- params = []
- for param in old_params:
- for model in model_types:
- if model not in nest.flatten(exclude_models):
- new_param = param.copy()
- new_param['_model_type'] = model
- new_param['testcase_name'] = '{}_{}'.format(
- param['testcase_name'], model)
- params.append(new_param)
-
- if extra_parameters:
- old_params = params
- params = []
- for param in old_params:
- for i, extra_params in enumerate(extra_parameters):
- new_param = param.copy()
- new_param.update(extra_params)
- new_param['testcase_name'] = '{}_{}'.format(
- param['testcase_name'], i)
- params.append(new_param)
-
- def decorator(f):
+ def single_method_decorator(f):
"""Decorator that constructs the test cases."""
- def run_with_model_type(self, model_type, *args, **kwargs):
- if model_type:
- with testing_utils.model_type_scope(model_type):
- f(self, *args, **kwargs)
- else:
- f(self, *args, **kwargs)
# Use named_parameters so it can be individually run from the command line
- @parameterized.named_parameters(
- *params
- )
- def decorated(self, *args, **kwargs):
- """A run of a single test case w/ specified run mode / model type."""
- run_mode = kwargs.pop('_run_mode', None)
- model_type = kwargs.pop('_model_type', None)
- if run_mode is None:
- run_with_model_type(self, model_type, *args, **kwargs)
- elif run_mode == 'v1_graph':
+ @parameterized.named_parameters(*params)
+ @functools.wraps(f)
+ def decorated(self, run_mode, *args, **kwargs):
+ """A run of a single test case w/ specified run mode."""
+ if run_mode == 'v1_graph':
with context.graph_mode(), testing_utils.run_eagerly_scope(False):
with self.test_session(use_gpu=True, config=config):
- run_with_model_type(self, model_type, *args, **kwargs)
+ f(self, *args, **kwargs)
elif run_mode == 'v2_function':
with context.eager_mode():
with testing_utils.run_eagerly_scope(False):
- run_with_model_type(self, model_type, *args, **kwargs)
+ f(self, *args, **kwargs)
elif run_mode == 'v2_eager':
with context.eager_mode():
with testing_utils.run_eagerly_scope(True):
- run_with_model_type(self, model_type, *args, **kwargs)
+ f(self, *args, **kwargs)
else:
return ValueError('Unknown run mode %s' % run_mode)
return decorated
- if func is not None:
- return decorator(func)
+ return _test_or_class_decorator(test_or_class, single_method_decorator)
- return decorator
+
+def _test_or_class_decorator(test_or_class, single_method_decorator):
+ """Decorate a test or class with a decorator intended for one method.
+
+ If the test_or_class is a class:
+ This will apply the decorator to all test methods in the class.
+
+ If the test_or_class is an iterable of already-parameterized test cases:
+ This will apply the decorator to all the cases, and then flatten the
+ resulting cross-product of test cases. This allows stacking the Keras
+ parameterized decorators w/ each other, and to apply them to test methods
+ that have already been marked with an absl parameterized decorator.
+
+ Otherwise, treat the obj as a single method and apply the decorator directly.
+
+ Args:
+ test_or_class: A test method (that may have already been decorated with a
+ parameterized decorator, or a test class that extends
+ keras_parameterized.TestCase
+ single_method_decorator:
+ A parameterized decorator intended for a single test method.
+ Returns:
+ The decorated result.
+ """
+ def _decorate_test_or_class(obj):
+ if isinstance(obj, collections.Iterable):
+ return itertools.chain.from_iterable(
+ single_method_decorator(method) for method in obj)
+ if isinstance(obj, type):
+ cls = obj
+ for name, value in cls.__dict__.copy().items():
+ if callable(value) and name.startswith(
+ unittest.TestLoader.testMethodPrefix):
+ setattr(cls, name, single_method_decorator(value))
+
+ cls = type(cls).__new__(type(cls), cls.__name__, cls.__bases__,
+ cls.__dict__.copy())
+ return cls
+
+ return single_method_decorator(obj)
+
+ if test_or_class is not None:
+ return _decorate_test_or_class(test_or_class)
+
+ return _decorate_test_or_class
diff --git a/tensorflow/python/keras/keras_parameterized_test.py b/tensorflow/python/keras/keras_parameterized_test.py
index 104c3b1..a0ff40c 100644
--- a/tensorflow/python/keras/keras_parameterized_test.py
+++ b/tensorflow/python/keras/keras_parameterized_test.py
@@ -20,6 +20,8 @@
import unittest
+from absl.testing import parameterized
+
from tensorflow.python import keras
from tensorflow.python import tf2
from tensorflow.python.eager import context
@@ -78,21 +80,22 @@
def runTest(self):
pass
- @keras_parameterized.run_with_all_model_types(
- extra_parameters=[dict(with_brackets=True), dict(with_brackets=False)]
- )
+ @keras_parameterized.run_with_all_model_types
+ @parameterized.named_parameters(
+ [dict(testcase_name="_0", with_brackets=True),
+ dict(testcase_name="_1", with_brackets=False)])
def testBody(self, with_brackets):
with_brackets = "with_brackets" if with_brackets else "without_brackets"
model_types.append((with_brackets, testing_utils.get_model_type()))
models.append(testing_utils.get_small_mlp(1, 4, input_dim=3))
e = ExampleTest()
- e.testBody_functional_0()
- e.testBody_subclass_0()
- e.testBody_sequential_0()
- e.testBody_functional_1()
- e.testBody_subclass_1()
- e.testBody_sequential_1()
+ e.testBody_0_functional()
+ e.testBody_0_subclass()
+ e.testBody_0_sequential()
+ e.testBody_1_functional()
+ e.testBody_1_subclass()
+ e.testBody_1_sequential()
self.assertLen(model_types, 6)
self.assertAllEqual(model_types, [
@@ -247,9 +250,10 @@
def runTest(self):
pass
- @keras_parameterized.run_all_keras_modes(
- extra_parameters=[dict(with_brackets=True), dict(with_brackets=False)]
- )
+ @keras_parameterized.run_all_keras_modes
+ @parameterized.named_parameters(
+ [dict(testcase_name="_0", with_brackets=True),
+ dict(testcase_name="_1", with_brackets=False)])
def testBody(self, with_brackets):
mode = "eager" if context.executing_eagerly() else "graph"
with_brackets = "with_brackets" if with_brackets else "without_brackets"
@@ -258,13 +262,13 @@
e = ExampleTest()
if not tf2.enabled():
- e.testBody_v1_graph_0()
- e.testBody_v1_graph_1()
+ e.testBody_0_v1_graph()
+ e.testBody_1_v1_graph()
- e.testBody_v2_eager_0()
- e.testBody_v2_function_0()
- e.testBody_v2_eager_1()
- e.testBody_v2_function_1()
+ e.testBody_0_v2_eager()
+ e.testBody_0_v2_function()
+ e.testBody_1_v2_eager()
+ e.testBody_1_v2_function()
expected_combinations = {
("with_brackets", "eager", True),
@@ -324,7 +328,8 @@
def runTest(self):
pass
- @keras_parameterized.run_all_keras_modes_with_all_model_types
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
def testBody(self):
mode = "eager" if context.executing_eagerly() else "graph"
should_run_eagerly = testing_utils.should_run_eagerly()
@@ -368,12 +373,178 @@
self.assertLen(l, len(expected_combinations) * 2)
- @keras_parameterized.run_all_keras_modes(extra_parameters=[dict(arg=True)])
+ def test_run_all_model_types_with_all_keras_modes(self):
+ l = []
+
+ class ExampleTest(keras_parameterized.TestCase):
+
+ def runTest(self):
+ pass
+
+ @keras_parameterized.run_all_keras_modes
+ @keras_parameterized.run_with_all_model_types
+ def testBody(self):
+ mode = "eager" if context.executing_eagerly() else "graph"
+ should_run_eagerly = testing_utils.should_run_eagerly()
+ l.append((mode, should_run_eagerly, testing_utils.get_model_type()))
+
+ e = ExampleTest()
+ e.testBody_functional_v2_eager()
+ e.testBody_functional_v2_function()
+ e.testBody_sequential_v2_eager()
+ e.testBody_sequential_v2_function()
+ e.testBody_subclass_v2_eager()
+ e.testBody_subclass_v2_function()
+
+ if not tf2.enabled():
+ e.testBody_functional_v1_graph()
+ e.testBody_sequential_v1_graph()
+ e.testBody_subclass_v1_graph()
+
+ expected_combinations = {
+ ("eager", True, "functional"),
+ ("eager", False, "functional"),
+ ("eager", True, "sequential"),
+ ("eager", False, "sequential"),
+ ("eager", True, "subclass"),
+ ("eager", False, "subclass"),
+ }
+
+ if not tf2.enabled():
+ expected_combinations = expected_combinations.union({
+ ("graph", False, "functional"),
+ ("graph", False, "sequential"),
+ ("graph", False, "subclass"),
+ })
+
+ self.assertLen(l, len(expected_combinations))
+ self.assertEqual(set(l), expected_combinations)
+
+ ts = unittest.makeSuite(ExampleTest)
+ res = unittest.TestResult()
+ ts.run(res)
+
+ self.assertLen(l, len(expected_combinations) * 2)
+
+ def test_run_all_keras_modes_with_all_model_types_annotate_class(self):
+ l = []
+
+ @keras_parameterized.run_with_all_model_types
+ @keras_parameterized.run_all_keras_modes
+ class ExampleTest(keras_parameterized.TestCase):
+
+ def runTest(self):
+ pass
+
+ @parameterized.named_parameters(dict(testcase_name="_arg",
+ arg=True))
+ def testBody(self, arg):
+ mode = "eager" if context.executing_eagerly() else "graph"
+ should_run_eagerly = testing_utils.should_run_eagerly()
+ l.append((mode, should_run_eagerly, testing_utils.get_model_type()))
+
+ e = ExampleTest()
+ e.testBody_arg_v2_eager_functional()
+ e.testBody_arg_v2_function_functional()
+ e.testBody_arg_v2_eager_sequential()
+ e.testBody_arg_v2_function_sequential()
+ e.testBody_arg_v2_eager_subclass()
+ e.testBody_arg_v2_function_subclass()
+
+ if not tf2.enabled():
+ e.testBody_arg_v1_graph_functional()
+ e.testBody_arg_v1_graph_sequential()
+ e.testBody_arg_v1_graph_subclass()
+
+ expected_combinations = {
+ ("eager", True, "functional"),
+ ("eager", False, "functional"),
+ ("eager", True, "sequential"),
+ ("eager", False, "sequential"),
+ ("eager", True, "subclass"),
+ ("eager", False, "subclass"),
+ }
+
+ if not tf2.enabled():
+ expected_combinations = expected_combinations.union({
+ ("graph", False, "functional"),
+ ("graph", False, "sequential"),
+ ("graph", False, "subclass"),
+ })
+
+ self.assertLen(l, len(expected_combinations))
+ self.assertEqual(set(l), expected_combinations)
+
+ ts = unittest.makeSuite(ExampleTest)
+ res = unittest.TestResult()
+ ts.run(res)
+
+ self.assertLen(l, len(expected_combinations) * 2)
+
+ def test_run_all_keras_modes_with_all_model_types_annotate_class_2(self):
+ l = []
+
+ @keras_parameterized.run_with_all_model_types
+ class ExampleTest(keras_parameterized.TestCase):
+
+ def runTest(self):
+ pass
+
+ @keras_parameterized.run_all_keras_modes
+ @parameterized.named_parameters(dict(testcase_name="_arg",
+ arg=True))
+ def testBody(self, arg):
+ mode = "eager" if context.executing_eagerly() else "graph"
+ should_run_eagerly = testing_utils.should_run_eagerly()
+ l.append((mode, should_run_eagerly, testing_utils.get_model_type()))
+
+ e = ExampleTest()
+ e.testBody_arg_v2_eager_functional()
+ e.testBody_arg_v2_function_functional()
+ e.testBody_arg_v2_eager_sequential()
+ e.testBody_arg_v2_function_sequential()
+ e.testBody_arg_v2_eager_subclass()
+ e.testBody_arg_v2_function_subclass()
+
+ if not tf2.enabled():
+ e.testBody_arg_v1_graph_functional()
+ e.testBody_arg_v1_graph_sequential()
+ e.testBody_arg_v1_graph_subclass()
+
+ expected_combinations = {
+ ("eager", True, "functional"),
+ ("eager", False, "functional"),
+ ("eager", True, "sequential"),
+ ("eager", False, "sequential"),
+ ("eager", True, "subclass"),
+ ("eager", False, "subclass"),
+ }
+
+ if not tf2.enabled():
+ expected_combinations = expected_combinations.union({
+ ("graph", False, "functional"),
+ ("graph", False, "sequential"),
+ ("graph", False, "subclass"),
+ })
+
+ self.assertLen(l, len(expected_combinations))
+ self.assertEqual(set(l), expected_combinations)
+
+ ts = unittest.makeSuite(ExampleTest)
+ res = unittest.TestResult()
+ ts.run(res)
+
+ self.assertLen(l, len(expected_combinations) * 2)
+
+ @keras_parameterized.run_all_keras_modes
+ @parameterized.named_parameters(dict(testcase_name="argument",
+ arg=True))
def test_run_all_keras_modes_extra_params_2(self, arg):
self.assertEqual(arg, True)
- @keras_parameterized.run_with_all_model_types(
- extra_parameters=[dict(arg=True)])
+ @keras_parameterized.run_with_all_model_types
+ @parameterized.named_parameters(dict(testcase_name="argument",
+ arg=True))
def test_run_with_all_model_types_extra_params_2(self, arg):
self.assertEqual(arg, True)
diff --git a/tensorflow/python/keras/layers/__init__.py b/tensorflow/python/keras/layers/__init__.py
index 49990b6b..df7571e 100644
--- a/tensorflow/python/keras/layers/__init__.py
+++ b/tensorflow/python/keras/layers/__init__.py
@@ -149,6 +149,7 @@
from tensorflow.python.keras.layers.recurrent import SimpleRNN
from tensorflow.python.keras.layers.recurrent import GRU
from tensorflow.python.keras.layers.recurrent import LSTM
+from tensorflow.python.keras.layers.recurrent import UnifiedLSTM
# Convolutional-recurrent layers.
from tensorflow.python.keras.layers.convolutional_recurrent import ConvLSTM2D
diff --git a/tensorflow/python/keras/layers/advanced_activations_test.py b/tensorflow/python/keras/layers/advanced_activations_test.py
index 4aadf53..f32bb45 100644
--- a/tensorflow/python/keras/layers/advanced_activations_test.py
+++ b/tensorflow/python/keras/layers/advanced_activations_test.py
@@ -20,13 +20,13 @@
from tensorflow.python import keras
from tensorflow.python.eager import context
-from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.platform import test
-@tf_test_util.run_all_in_graph_and_eager_modes
-class AdvancedActivationsTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class AdvancedActivationsTest(keras_parameterized.TestCase):
def test_leaky_relu(self):
for alpha in [0., .5, -1.]:
diff --git a/tensorflow/python/keras/layers/convolutional_test.py b/tensorflow/python/keras/layers/convolutional_test.py
index d3339a8..81af06b 100644
--- a/tensorflow/python/keras/layers/convolutional_test.py
+++ b/tensorflow/python/keras/layers/convolutional_test.py
@@ -24,13 +24,13 @@
from tensorflow.python import keras
from tensorflow.python.eager import context
-from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.platform import test
-@tf_test_util.run_all_in_graph_and_eager_modes
-class Convolution1DTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class Convolution1DTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -100,8 +100,8 @@
self.assertEqual(layer.bias.constraint, b_constraint)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class Conv2DTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class Conv2DTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -175,8 +175,8 @@
self.assertEqual(layer.bias.constraint, b_constraint)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class Conv2DTransposeTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class Conv2DTransposeTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -267,8 +267,8 @@
expected_output=expected_output)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class Conv3DTransposeTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class Conv3DTransposeTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -336,8 +336,8 @@
self.assertEqual(layer.bias.constraint, b_constraint)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class SeparableConv1DTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class SeparableConv1DTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -411,8 +411,8 @@
self.assertEqual(layer.bias.constraint, b_constraint)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class SeparableConv2DTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class SeparableConv2DTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -489,8 +489,8 @@
self.assertEqual(layer.bias.constraint, b_constraint)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class Conv3DTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class Conv3DTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -557,8 +557,8 @@
self.assertEqual(layer.bias.constraint, b_constraint)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class ZeroPaddingTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class ZeroPaddingTest(keras_parameterized.TestCase):
def test_zero_padding_1d(self):
num_samples = 2
@@ -726,8 +726,8 @@
keras.layers.ZeroPadding3D(padding=None)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class UpSamplingTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class UpSamplingTest(keras_parameterized.TestCase):
def test_upsampling_1d(self):
with self.session(use_gpu=True):
@@ -875,8 +875,8 @@
np.testing.assert_allclose(np_output, expected_out)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class CroppingTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class CroppingTest(keras_parameterized.TestCase):
def test_cropping_1d(self):
num_samples = 2
@@ -1017,8 +1017,8 @@
keras.layers.Cropping3D(cropping=None)
-@tf_test_util.run_all_in_graph_and_eager_modes
-class DepthwiseConv2DTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class DepthwiseConv2DTest(keras_parameterized.TestCase):
def _run_test(self, kwargs, arg, values):
num_samples = 2
@@ -1044,17 +1044,18 @@
self._run_test(kwargs, 'data_format', ['channels_first'])
self._run_test(kwargs, 'depth_multiplier', [1, 2])
- kwargs = {'kernel_size': 3,
- 'padding': 'valid',
- 'data_format': 'channels_first',
- 'activation': None,
- 'depthwise_regularizer': 'l2',
- 'bias_regularizer': 'l2',
- 'activity_regularizer': 'l2',
- 'depthwise_constraint': 'unit_norm',
- 'use_bias': True,
- 'strides': (2, 2),
- }
+ kwargs = {
+ 'kernel_size': 3,
+ 'padding': 'valid',
+ 'data_format': 'channels_last',
+ 'activation': None,
+ 'depthwise_regularizer': 'l2',
+ 'bias_regularizer': 'l2',
+ 'activity_regularizer': 'l2',
+ 'depthwise_constraint': 'unit_norm',
+ 'use_bias': True,
+ 'strides': (2, 2),
+ }
self._run_test(kwargs, 'depth_multiplier', [1])
if __name__ == '__main__':
diff --git a/tensorflow/python/keras/layers/core.py b/tensorflow/python/keras/layers/core.py
index 56dd705..39bcb82 100644
--- a/tensorflow/python/keras/layers/core.py
+++ b/tensorflow/python/keras/layers/core.py
@@ -506,6 +506,9 @@
class Flatten(Layer):
"""Flattens the input. Does not affect the batch size.
+ If inputs are shaped `(batch,)` without a channel dimension, then flattening
+ adds an extra channel dimension and output shapes are `(batch, 1)`.
+
Arguments:
data_format: A string,
one of `channels_last` (default) or `channels_first`.
@@ -534,23 +537,28 @@
def __init__(self, data_format=None, **kwargs):
super(Flatten, self).__init__(**kwargs)
self.data_format = conv_utils.normalize_data_format(data_format)
- self.input_spec = InputSpec(min_ndim=2)
+ self.input_spec = InputSpec(min_ndim=1)
def call(self, inputs):
- if self.data_format == 'channels_first':
+ if (self.data_format == 'channels_first'
+ and K.ndim(inputs) is not None and K.ndim(inputs) > 1):
permutation = [0]
permutation.extend([i for i in
range(2, K.ndim(inputs))])
permutation.append(1)
inputs = array_ops.transpose(inputs, perm=permutation)
- outputs = array_ops.reshape(inputs, (array_ops.shape(inputs)[0], -1))
+ outputs = array_ops.reshape(
+ inputs, (tensor_shape.dimension_value(inputs.shape[0]) or
+ array_ops.shape(inputs)[0], -1))
if not context.executing_eagerly():
outputs.set_shape(self.compute_output_shape(inputs.get_shape()))
return outputs
def compute_output_shape(self, input_shape):
input_shape = tensor_shape.TensorShape(input_shape).as_list()
+ if not input_shape:
+ output_shape = tensor_shape.TensorShape([1])
output_shape = [input_shape[0]]
if all(input_shape[1:]):
output_shape += [np.prod(input_shape[1:])]
diff --git a/tensorflow/python/keras/layers/core_test.py b/tensorflow/python/keras/layers/core_test.py
index aad6ab8..9df40f8 100644
--- a/tensorflow/python/keras/layers/core_test.py
+++ b/tensorflow/python/keras/layers/core_test.py
@@ -22,43 +22,36 @@
from tensorflow.python import keras
from tensorflow.python.eager import context
-from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
-class CoreLayersTest(test.TestCase):
-
- def test_masking(self):
- with self.cached_session():
- testing_utils.layer_test(
- keras.layers.Masking, kwargs={}, input_shape=(3, 2, 3))
+@keras_parameterized.run_all_keras_modes
+class DropoutLayersTest(keras_parameterized.TestCase):
def test_dropout(self):
- with self.cached_session():
- testing_utils.layer_test(
- keras.layers.Dropout, kwargs={'rate': 0.5}, input_shape=(3, 2))
+ testing_utils.layer_test(
+ keras.layers.Dropout, kwargs={'rate': 0.5}, input_shape=(3, 2))
- with self.cached_session():
- testing_utils.layer_test(
- keras.layers.Dropout,
- kwargs={'rate': 0.5,
- 'noise_shape': [3, 1]},
- input_shape=(3, 2))
+ testing_utils.layer_test(
+ keras.layers.Dropout,
+ kwargs={'rate': 0.5,
+ 'noise_shape': [3, 1]},
+ input_shape=(3, 2))
- # https://github.com/tensorflow/tensorflow/issues/14819
- with self.cached_session():
- dropout = keras.layers.Dropout(0.5)
- self.assertEqual(True, dropout.supports_masking)
+ def test_dropout_supports_masking(self):
+ dropout = keras.layers.Dropout(0.5)
+ self.assertEqual(True, dropout.supports_masking)
- @tf_test_util.run_in_graph_and_eager_modes
- def test_spatial_dropout(self):
+ def test_spatial_dropout_1d(self):
testing_utils.layer_test(
keras.layers.SpatialDropout1D,
kwargs={'rate': 0.5},
input_shape=(2, 3, 4))
+ def test_spatial_dropout_2d(self):
testing_utils.layer_test(
keras.layers.SpatialDropout2D,
kwargs={'rate': 0.5},
@@ -69,6 +62,7 @@
kwargs={'rate': 0.5, 'data_format': 'channels_first'},
input_shape=(2, 3, 4, 5))
+ def test_spatial_dropout_3d(self):
testing_utils.layer_test(
keras.layers.SpatialDropout3D,
kwargs={'rate': 0.5},
@@ -79,80 +73,9 @@
kwargs={'rate': 0.5, 'data_format': 'channels_first'},
input_shape=(2, 3, 4, 4, 5))
- @tf_test_util.run_in_graph_and_eager_modes
- def test_activation(self):
- # with string argument
- testing_utils.layer_test(
- keras.layers.Activation,
- kwargs={'activation': 'relu'},
- input_shape=(3, 2))
- # with function argument
- testing_utils.layer_test(
- keras.layers.Activation,
- kwargs={'activation': keras.backend.relu},
- input_shape=(3, 2))
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_reshape(self):
- testing_utils.layer_test(
- keras.layers.Reshape,
- kwargs={'target_shape': (8, 1)},
- input_shape=(3, 2, 4))
-
- testing_utils.layer_test(
- keras.layers.Reshape,
- kwargs={'target_shape': (-1, 1)},
- input_shape=(3, 2, 4))
-
- testing_utils.layer_test(
- keras.layers.Reshape,
- kwargs={'target_shape': (1, -1)},
- input_shape=(3, 2, 4))
-
- testing_utils.layer_test(
- keras.layers.Reshape,
- kwargs={'target_shape': (-1, 1)},
- input_shape=(None, None, 2))
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_permute(self):
- testing_utils.layer_test(
- keras.layers.Permute, kwargs={'dims': (2, 1)}, input_shape=(3, 2, 4))
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_permute_errors_on_invalid_starting_dims_index(self):
- with self.assertRaisesRegexp(ValueError, r'Invalid permutation .*dims.*'):
- testing_utils.layer_test(
- keras.layers.Permute,
- kwargs={'dims': (0, 1, 2)}, input_shape=(3, 2, 4))
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_permute_errors_on_invalid_set_of_dims_indices(self):
- with self.assertRaisesRegexp(ValueError, r'Invalid permutation .*dims.*'):
- testing_utils.layer_test(
- keras.layers.Permute,
- kwargs={'dims': (1, 4, 2)}, input_shape=(3, 2, 4))
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_flatten(self):
- testing_utils.layer_test(
- keras.layers.Flatten, kwargs={}, input_shape=(3, 2, 4))
-
- # Test channels_first
- inputs = np.random.random((10, 3, 5, 5)).astype('float32')
- outputs = testing_utils.layer_test(
- keras.layers.Flatten,
- kwargs={'data_format': 'channels_first'},
- input_data=inputs)
- target_outputs = np.reshape(
- np.transpose(inputs, (0, 2, 3, 1)), (-1, 5 * 5 * 3))
- self.assertAllClose(outputs, target_outputs)
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_repeat_vector(self):
- testing_utils.layer_test(
- keras.layers.RepeatVector, kwargs={'n': 3}, input_shape=(3, 2))
+@keras_parameterized.run_all_keras_modes
+class LambdaLayerTest(keras_parameterized.TestCase):
def test_lambda(self):
testing_utils.layer_test(
@@ -188,7 +111,6 @@
config = ld.get_config()
ld = keras.layers.Lambda.from_config(config)
- @tf_test_util.run_in_graph_and_eager_modes
def test_lambda_multiple_inputs(self):
ld = keras.layers.Lambda(lambda x: x[0], output_shape=lambda x: x[0])
x1 = np.ones([3, 2], np.float32)
@@ -196,7 +118,158 @@
out = ld([x1, x2])
self.assertAllEqual(out.shape, [3, 2])
- @tf_test_util.run_in_graph_and_eager_modes
+ def test_lambda_output_shape(self):
+ l = keras.layers.Lambda(lambda x: x + 1, output_shape=(1, 1))
+ l(keras.backend.variable(np.ones((1, 1))))
+ self.assertEqual((1, 1), l.get_config()['output_shape'])
+
+ def test_lambda_output_shape_function(self):
+ def get_output_shape(input_shape):
+ return 1 * input_shape
+
+ l = keras.layers.Lambda(lambda x: x + 1, output_shape=get_output_shape)
+ l(keras.backend.variable(np.ones((1, 1))))
+ self.assertEqual('lambda', l.get_config()['output_shape_type'])
+
+ def test_lambda_output_shape_autocalculate_multiple_inputs(self):
+
+ def lambda_fn(x):
+ return math_ops.matmul(x[0], x[1])
+
+ l = keras.layers.Lambda(lambda_fn)
+ output_shape = l.compute_output_shape([(10, 10), (10, 20)])
+ self.assertAllEqual((10, 20), output_shape)
+
+ def test_lambda_output_shape_list_multiple_outputs(self):
+
+ def lambda_fn(x):
+ return x
+
+ l = keras.layers.Lambda(lambda_fn, output_shape=[(10,), (20,)])
+ output_shape = l.compute_output_shape([(10, 10), (10, 20)])
+ self.assertAllEqual([(10, 10), (10, 20)], output_shape)
+
+ def test_lambda_output_shape_tuple_with_none(self):
+
+ def lambda_fn(x):
+ return x
+
+ l = keras.layers.Lambda(lambda_fn, output_shape=(None, 10))
+ output_shape = l.compute_output_shape((5, 10, 20))
+ self.assertAllEqual([5, None, 10], output_shape.as_list())
+
+ def test_lambda_output_shape_function_multiple_outputs(self):
+
+ def lambda_fn(x):
+ return x
+
+ def output_shape_fn(input_shape):
+ return input_shape
+
+ l = keras.layers.Lambda(lambda_fn, output_shape=output_shape_fn)
+ output_shape = l.compute_output_shape([(10, 10), (10, 20)])
+ self.assertAllEqual([(10, 10), (10, 20)], output_shape)
+
+ def test_lambda_config_serialization(self):
+ # Test serialization with output_shape and output_shape_type
+ layer = keras.layers.Lambda(lambda x: x + 1, output_shape=(1, 1))
+ layer(keras.backend.variable(np.ones((1, 1))))
+ config = layer.get_config()
+ layer = keras.layers.deserialize({
+ 'class_name': 'Lambda',
+ 'config': config
+ })
+ layer = keras.layers.Lambda.from_config(config)
+
+
+@keras_parameterized.run_all_keras_modes
+class CoreLayersTest(keras_parameterized.TestCase):
+
+ def test_masking(self):
+ testing_utils.layer_test(
+ keras.layers.Masking, kwargs={}, input_shape=(3, 2, 3))
+
+ def test_activation(self):
+ # with string argument
+ testing_utils.layer_test(
+ keras.layers.Activation,
+ kwargs={'activation': 'relu'},
+ input_shape=(3, 2))
+
+ # with function argument
+ testing_utils.layer_test(
+ keras.layers.Activation,
+ kwargs={'activation': keras.backend.relu},
+ input_shape=(3, 2))
+
+ def test_reshape(self):
+ testing_utils.layer_test(
+ keras.layers.Reshape,
+ kwargs={'target_shape': (8, 1)},
+ input_shape=(3, 2, 4))
+
+ testing_utils.layer_test(
+ keras.layers.Reshape,
+ kwargs={'target_shape': (-1, 1)},
+ input_shape=(3, 2, 4))
+
+ testing_utils.layer_test(
+ keras.layers.Reshape,
+ kwargs={'target_shape': (1, -1)},
+ input_shape=(3, 2, 4))
+
+ testing_utils.layer_test(
+ keras.layers.Reshape,
+ kwargs={'target_shape': (-1, 1)},
+ input_shape=(None, None, 2))
+
+ def test_permute(self):
+ testing_utils.layer_test(
+ keras.layers.Permute, kwargs={'dims': (2, 1)}, input_shape=(3, 2, 4))
+
+ def test_permute_errors_on_invalid_starting_dims_index(self):
+ with self.assertRaisesRegexp(ValueError, r'Invalid permutation .*dims.*'):
+ testing_utils.layer_test(
+ keras.layers.Permute,
+ kwargs={'dims': (0, 1, 2)}, input_shape=(3, 2, 4))
+
+ def test_permute_errors_on_invalid_set_of_dims_indices(self):
+ with self.assertRaisesRegexp(ValueError, r'Invalid permutation .*dims.*'):
+ testing_utils.layer_test(
+ keras.layers.Permute,
+ kwargs={'dims': (1, 4, 2)}, input_shape=(3, 2, 4))
+
+ def test_flatten(self):
+ testing_utils.layer_test(
+ keras.layers.Flatten, kwargs={}, input_shape=(3, 2, 4))
+
+ # Test channels_first
+ inputs = np.random.random((10, 3, 5, 5)).astype('float32')
+ outputs = testing_utils.layer_test(
+ keras.layers.Flatten,
+ kwargs={'data_format': 'channels_first'},
+ input_data=inputs)
+ target_outputs = np.reshape(
+ np.transpose(inputs, (0, 2, 3, 1)), (-1, 5 * 5 * 3))
+ self.assertAllClose(outputs, target_outputs)
+
+ def test_flatten_scalar_channels(self):
+ testing_utils.layer_test(
+ keras.layers.Flatten, kwargs={}, input_shape=(3,))
+
+ # Test channels_first
+ inputs = np.random.random((10,)).astype('float32')
+ outputs = testing_utils.layer_test(
+ keras.layers.Flatten,
+ kwargs={'data_format': 'channels_first'},
+ input_data=inputs)
+ target_outputs = np.expand_dims(inputs, -1)
+ self.assertAllClose(outputs, target_outputs)
+
+ def test_repeat_vector(self):
+ testing_utils.layer_test(
+ keras.layers.RepeatVector, kwargs={'n': 3}, input_shape=(3, 2))
+
def test_dense(self):
testing_utils.layer_test(
keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 2))
@@ -211,105 +284,31 @@
keras.layers.Dense, kwargs={'units': 3}, input_shape=(3, 4, 5, 2))
def test_dense_regularization(self):
- with self.cached_session():
- layer = keras.layers.Dense(
- 3,
- kernel_regularizer=keras.regularizers.l1(0.01),
- bias_regularizer='l1',
- activity_regularizer='l2',
- name='dense_reg')
- layer(keras.backend.variable(np.ones((2, 4))))
- self.assertEqual(3, len(layer.losses))
+ layer = keras.layers.Dense(
+ 3,
+ kernel_regularizer=keras.regularizers.l1(0.01),
+ bias_regularizer='l1',
+ activity_regularizer='l2',
+ name='dense_reg')
+ layer(keras.backend.variable(np.ones((2, 4))))
+ self.assertEqual(3, len(layer.losses))
def test_dense_constraints(self):
- with self.cached_session():
- k_constraint = keras.constraints.max_norm(0.01)
- b_constraint = keras.constraints.max_norm(0.01)
- layer = keras.layers.Dense(
- 3, kernel_constraint=k_constraint, bias_constraint=b_constraint)
- layer(keras.backend.variable(np.ones((2, 4))))
- self.assertEqual(layer.kernel.constraint, k_constraint)
- self.assertEqual(layer.bias.constraint, b_constraint)
+ k_constraint = keras.constraints.max_norm(0.01)
+ b_constraint = keras.constraints.max_norm(0.01)
+ layer = keras.layers.Dense(
+ 3, kernel_constraint=k_constraint, bias_constraint=b_constraint)
+ layer(keras.backend.variable(np.ones((2, 4))))
+ self.assertEqual(layer.kernel.constraint, k_constraint)
+ self.assertEqual(layer.bias.constraint, b_constraint)
def test_activity_regularization(self):
- with self.cached_session():
- layer = keras.layers.ActivityRegularization(l1=0.1)
- layer(keras.backend.variable(np.ones((2, 4))))
- self.assertEqual(1, len(layer.losses))
- _ = layer.get_config()
+ layer = keras.layers.ActivityRegularization(l1=0.1)
+ layer(keras.backend.variable(np.ones((2, 4))))
+ self.assertEqual(1, len(layer.losses))
+ config = layer.get_config()
+ self.assertEqual(config.pop('l1'), 0.1)
- def test_lambda_output_shape(self):
- with self.cached_session():
- l = keras.layers.Lambda(lambda x: x + 1, output_shape=(1, 1))
- l(keras.backend.variable(np.ones((1, 1))))
- self.assertEqual((1, 1), l.get_config()['output_shape'])
-
- def test_lambda_output_shape_function(self):
- def get_output_shape(input_shape):
- return 1 * input_shape
-
- with self.cached_session():
- l = keras.layers.Lambda(lambda x: x + 1, output_shape=get_output_shape)
- l(keras.backend.variable(np.ones((1, 1))))
- self.assertEqual('lambda', l.get_config()['output_shape_type'])
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_lambda_output_shape_autocalculate_multiple_inputs(self):
-
- def lambda_fn(x):
- return math_ops.matmul(x[0], x[1])
-
- l = keras.layers.Lambda(lambda_fn)
- output_shape = l.compute_output_shape([(10, 10), (10, 20)])
- self.assertAllEqual((10, 20), output_shape)
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_lambda_output_shape_list_multiple_outputs(self):
-
- def lambda_fn(x):
- return x
-
- l = keras.layers.Lambda(lambda_fn, output_shape=[(10,), (20,)])
- output_shape = l.compute_output_shape([(10, 10), (10, 20)])
- self.assertAllEqual([(10, 10), (10, 20)], output_shape)
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_lambda_output_shape_tuple_with_none(self):
-
- def lambda_fn(x):
- return x
-
- l = keras.layers.Lambda(lambda_fn, output_shape=(None, 10))
- output_shape = l.compute_output_shape((5, 10, 20))
- self.assertAllEqual([5, None, 10], output_shape.as_list())
-
- @tf_test_util.run_in_graph_and_eager_modes
- def test_lambda_output_shape_function_multiple_outputs(self):
-
- def lambda_fn(x):
- return x
-
- def output_shape_fn(input_shape):
- return input_shape
-
- l = keras.layers.Lambda(lambda_fn, output_shape=output_shape_fn)
- output_shape = l.compute_output_shape([(10, 10), (10, 20)])
- self.assertAllEqual([(10, 10), (10, 20)], output_shape)
-
- def test_lambda_config_serialization(self):
- with self.cached_session():
- # test serialization with output_shape and output_shape_type
- layer = keras.layers.Lambda(lambda x: x + 1, output_shape=(1, 1))
- layer(keras.backend.variable(np.ones((1, 1))))
- config = layer.get_config()
- layer = keras.layers.deserialize({
- 'class_name': 'Lambda',
- 'config': config
- })
-
- layer = keras.layers.Lambda.from_config(config)
-
- @tf_test_util.run_in_graph_and_eager_modes
def test_numpy_inputs(self):
if context.executing_eagerly():
layer = keras.layers.RepeatVector(2)
diff --git a/tensorflow/python/keras/layers/cudnn_recurrent.py b/tensorflow/python/keras/layers/cudnn_recurrent.py
index 1669275..e695a68 100644
--- a/tensorflow/python/keras/layers/cudnn_recurrent.py
+++ b/tensorflow/python/keras/layers/cudnn_recurrent.py
@@ -335,7 +335,7 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('keras.layers.CuDNNLSTM')
+@tf_export(v1=['keras.layers.CuDNNLSTM'])
class CuDNNLSTM(_CuDNNRNN):
"""Fast LSTM implementation backed by cuDNN.
diff --git a/tensorflow/python/keras/layers/gru_test.py b/tensorflow/python/keras/layers/gru_test.py
index 9988c9f..1b2881a 100644
--- a/tensorflow/python/keras/layers/gru_test.py
+++ b/tensorflow/python/keras/layers/gru_test.py
@@ -81,6 +81,29 @@
'implementation': mode},
input_shape=(num_samples, timesteps, embedding_dim))
+ @tf_test_util.run_in_graph_and_eager_modes
+ def test_reset_after_GRU(self):
+ num_samples = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 2
+
+ (x_train, y_train), _ = testing_utils.get_test_data(
+ train_samples=num_samples,
+ test_samples=0,
+ input_shape=(timesteps, embedding_dim),
+ num_classes=units)
+ y_train = keras.utils.to_categorical(y_train, units)
+
+ inputs = keras.layers.Input(shape=[timesteps, embedding_dim])
+ gru_layer = keras.layers.GRU(units,
+ reset_after=True)
+ output = gru_layer(inputs)
+ gru_model = keras.models.Model(inputs, output)
+ gru_model.compile('rmsprop', 'mse')
+ gru_model.fit(x_train, y_train)
+ gru_model.predict(x_train)
+
def test_statefulness_GRU(self):
num_samples = 2
timesteps = 3
diff --git a/tensorflow/python/keras/layers/local_test.py b/tensorflow/python/keras/layers/local_test.py
index e4f4d0a..6db5bf3 100644
--- a/tensorflow/python/keras/layers/local_test.py
+++ b/tensorflow/python/keras/layers/local_test.py
@@ -235,7 +235,7 @@
class LocallyConnectedImplementationModeTest(test.TestCase):
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_locallyconnected_implementation(self):
with self.cached_session():
num_samples = 4
diff --git a/tensorflow/python/keras/layers/lstm_test.py b/tensorflow/python/keras/layers/lstm_test.py
index 3f89cc3..5c0ad5c 100644
--- a/tensorflow/python/keras/layers/lstm_test.py
+++ b/tensorflow/python/keras/layers/lstm_test.py
@@ -23,6 +23,7 @@
from tensorflow.python import keras
from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.platform import test
from tensorflow.python.training import adam
@@ -30,8 +31,8 @@
from tensorflow.python.training.rmsprop import RMSPropOptimizer
-@tf_test_util.run_all_in_graph_and_eager_modes
-class LSTMLayerTest(test.TestCase, parameterized.TestCase):
+@keras_parameterized.run_all_keras_modes
+class LSTMLayerTest(keras_parameterized.TestCase):
def test_return_sequences_LSTM(self):
num_samples = 2
@@ -67,7 +68,9 @@
layer = keras.layers.LSTM(units, input_shape=(None, embedding_dim))
model = keras.models.Sequential()
model.add(layer)
- model.compile(RMSPropOptimizer(0.001), 'mse')
+ model.compile(RMSPropOptimizer(0.001), 'mse',
+ run_eagerly=testing_utils.should_run_eagerly())
+
x = np.random.random((num_samples, timesteps, embedding_dim))
y = np.random.random((num_samples, units))
model.train_on_batch(x, y)
@@ -115,7 +118,6 @@
self.assertEqual(layer.cell.recurrent_kernel.constraint, r_constraint)
self.assertEqual(layer.cell.bias.constraint, b_constraint)
- @tf_test_util.run_deprecated_v1
def test_with_masking_layer_LSTM(self):
layer_class = keras.layers.LSTM
inputs = np.random.random((2, 3, 4))
@@ -125,10 +127,10 @@
model.add(keras.layers.Masking(input_shape=(3, 4)))
model.add(layer_class(units=5, return_sequences=True, unroll=False))
model.compile(loss='categorical_crossentropy',
- optimizer=RMSPropOptimizer(0.01))
+ optimizer=RMSPropOptimizer(0.01),
+ run_eagerly=testing_utils.should_run_eagerly())
model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
- @tf_test_util.run_deprecated_v1
def test_masking_with_stacking_LSTM(self):
inputs = np.random.random((2, 3, 4))
targets = np.abs(np.random.random((2, 3, 5)))
@@ -138,7 +140,8 @@
lstm_cells = [keras.layers.LSTMCell(10), keras.layers.LSTMCell(5)]
model.add(keras.layers.RNN(lstm_cells, return_sequences=True, unroll=False))
model.compile(loss='categorical_crossentropy',
- optimizer=RMSPropOptimizer(0.01))
+ optimizer=RMSPropOptimizer(0.01),
+ run_eagerly=testing_utils.should_run_eagerly())
model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
def test_from_config_LSTM(self):
@@ -167,7 +170,8 @@
model = keras.models.Model([inputs] + initial_state, output)
model.compile(loss='categorical_crossentropy',
- optimizer=adam.AdamOptimizer())
+ optimizer=adam.AdamOptimizer(),
+ run_eagerly=testing_utils.should_run_eagerly())
inputs = np.random.random((num_samples, timesteps, embedding_dim))
initial_state = [np.random.random((num_samples, units))
@@ -192,7 +196,8 @@
model = keras.models.Model(inputs, output)
model.compile(loss='categorical_crossentropy',
- optimizer=adam.AdamOptimizer())
+ optimizer=adam.AdamOptimizer(),
+ run_eagerly=testing_utils.should_run_eagerly())
inputs = np.random.random((num_samples, timesteps, embedding_dim))
targets = np.random.random((num_samples, units))
@@ -242,7 +247,8 @@
model = keras.models.Model([inputs] + initial_state, output)
model.compile(loss='categorical_crossentropy',
- optimizer=RMSPropOptimizer(0.01))
+ optimizer=RMSPropOptimizer(0.01),
+ run_eagerly=testing_utils.should_run_eagerly())
inputs = np.random.random((num_samples, timesteps, embedding_dim))
initial_state = [np.random.random((num_samples, units))
@@ -303,7 +309,8 @@
model = keras.models.Model(inputs, output)
model.compile(loss='categorical_crossentropy',
- optimizer=adam.AdamOptimizer())
+ optimizer=adam.AdamOptimizer(),
+ run_eagerly=testing_utils.should_run_eagerly())
main_inputs = np.random.random((num_samples, timesteps, embedding_dim))
initial_state = [np.random.random((num_samples, units))
@@ -314,7 +321,7 @@
class LSTMLayerGraphOnlyTest(test.TestCase):
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_statefulness_LSTM(self):
num_samples = 2
timesteps = 3
@@ -378,25 +385,25 @@
self.assertAllClose(out7, out6, atol=1e-5)
+ # b/120919032
@tf_test_util.run_deprecated_v1
def test_regularizers_LSTM(self):
embedding_dim = 4
layer_class = keras.layers.LSTM
- with self.cached_session():
- layer = layer_class(
- 5,
- return_sequences=False,
- weights=None,
- input_shape=(None, embedding_dim),
- kernel_regularizer=keras.regularizers.l1(0.01),
- recurrent_regularizer=keras.regularizers.l1(0.01),
- bias_regularizer='l2',
- activity_regularizer='l1')
- layer.build((None, None, 2))
- self.assertEqual(len(layer.losses), 3)
- x = keras.backend.variable(np.ones((2, 3, 2)))
- layer(x)
- self.assertEqual(len(layer.get_losses_for(x)), 1)
+ layer = layer_class(
+ 5,
+ return_sequences=False,
+ weights=None,
+ input_shape=(None, embedding_dim),
+ kernel_regularizer=keras.regularizers.l1(0.01),
+ recurrent_regularizer=keras.regularizers.l1(0.01),
+ bias_regularizer='l2',
+ activity_regularizer='l1')
+ layer.build((None, None, 2))
+ self.assertEqual(len(layer.losses), 3)
+ x = keras.backend.variable(np.ones((2, 3, 2)))
+ layer(x)
+ self.assertEqual(len(layer.get_losses_for(x)), 1)
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/layers/merge.py b/tensorflow/python/keras/layers/merge.py
index 45e705c..c73b21d 100644
--- a/tensorflow/python/keras/layers/merge.py
+++ b/tensorflow/python/keras/layers/merge.py
@@ -87,7 +87,7 @@
def build(self, input_shape):
# Used purely for shape validation.
if not isinstance(input_shape, list):
- raise ValueError('A merge layer should be called ' 'on a list of inputs.')
+ raise ValueError('A merge layer should be called on a list of inputs.')
if len(input_shape) < 2:
raise ValueError('A merge layer should be called '
'on a list of at least 2 inputs. '
@@ -118,7 +118,7 @@
def call(self, inputs):
if not isinstance(inputs, list):
- raise ValueError('A merge layer should be called ' 'on a list of inputs.')
+ raise ValueError('A merge layer should be called on a list of inputs.')
if self._reshape_required:
reshaped_inputs = []
input_ndims = list(map(K.ndim, inputs))
@@ -504,7 +504,7 @@
def _merge_function(self, inputs):
if len(inputs) != 2:
- raise ValueError('A `Dot` layer should be called ' 'on exactly 2 inputs')
+ raise ValueError('A `Dot` layer should be called on exactly 2 inputs')
x1 = inputs[0]
x2 = inputs[1]
if isinstance(self.axes, int):
diff --git a/tensorflow/python/keras/layers/merge_test.py b/tensorflow/python/keras/layers/merge_test.py
index fcb161a..f962a75 100644
--- a/tensorflow/python/keras/layers/merge_test.py
+++ b/tensorflow/python/keras/layers/merge_test.py
@@ -22,12 +22,13 @@
from tensorflow.python import keras
from tensorflow.python.framework import test_util as tf_test_util
-from tensorflow.python.ops import array_ops
+from tensorflow.python.keras import keras_parameterized
+from tensorflow.python.keras import testing_utils
from tensorflow.python.platform import test
-@tf_test_util.run_all_in_graph_and_eager_modes
-class MergeLayersTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class MergeLayersTest(keras_parameterized.TestCase):
def test_merge_add(self):
i1 = keras.layers.Input(shape=(4, 5))
@@ -35,8 +36,9 @@
i3 = keras.layers.Input(shape=(4, 5))
o = keras.layers.add([i1, i2, i3])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 5])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 5])
model = keras.models.Model([i1, i2, i3], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
x1 = np.random.random((2, 4, 5))
x2 = np.random.random((2, 4, 5))
@@ -45,25 +47,14 @@
self.assertEqual(out.shape, (2, 4, 5))
self.assertAllClose(out, x1 + x2 + x3, atol=1e-4)
- def test_merge_elementwise_errors(self):
- i1 = keras.layers.Input(shape=(4, 5))
- i2 = keras.layers.Input(shape=(4, 6))
- with self.assertRaises(ValueError):
- keras.layers.add([i1, i2])
- with self.assertRaises(ValueError):
- keras.layers.add([i1])
- with self.assertRaises(ValueError):
- keras.layers.add(i1)
- with self.assertRaises(ValueError):
- keras.layers.add([i1])
-
def test_merge_multiply(self):
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 5))
i3 = keras.layers.Input(shape=(4, 5))
o = keras.layers.multiply([i1, i2, i3])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 5])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 5])
model = keras.models.Model([i1, i2, i3], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
x1 = np.random.random((2, 4, 5))
x2 = np.random.random((2, 4, 5))
@@ -76,8 +67,9 @@
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 5))
o = keras.layers.average([i1, i2])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 5])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 5])
model = keras.models.Model([i1, i2], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
x1 = np.random.random((2, 4, 5))
x2 = np.random.random((2, 4, 5))
@@ -89,8 +81,9 @@
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 5))
o = keras.layers.maximum([i1, i2])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 5])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 5])
model = keras.models.Model([i1, i2], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
x1 = np.random.random((2, 4, 5))
x2 = np.random.random((2, 4, 5))
@@ -102,8 +95,9 @@
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 5))
o = keras.layers.minimum([i1, i2])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 5])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 5])
model = keras.models.Model([i1, i2], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
x1 = np.random.random((2, 4, 5))
x2 = np.random.random((2, 4, 5))
@@ -115,8 +109,9 @@
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 5))
o = keras.layers.concatenate([i1, i2], axis=1)
- self.assertListEqual(o.get_shape().as_list(), [None, 8, 5])
+ self.assertListEqual(o.shape.as_list(), [None, 8, 5])
model = keras.models.Model([i1, i2], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
x1 = np.random.random((2, 4, 5))
x2 = np.random.random((2, 4, 5))
@@ -124,22 +119,13 @@
self.assertEqual(out.shape, (2, 8, 5))
self.assertAllClose(out, np.concatenate([x1, x2], axis=1), atol=1e-4)
- def test_concatenate_errors(self):
- i1 = keras.layers.Input(shape=(4, 5))
- i2 = keras.layers.Input(shape=(3, 5))
- with self.assertRaisesRegexp(ValueError, 'inputs with matching shapes'):
- keras.layers.concatenate([i1, i2], axis=-1)
- with self.assertRaisesRegexp(ValueError, 'called on a list'):
- keras.layers.concatenate(i1, axis=-1)
- with self.assertRaisesRegexp(ValueError, 'called on a list'):
- keras.layers.concatenate([i1], axis=-1)
-
def test_merge_dot(self):
i1 = keras.layers.Input(shape=(4,))
i2 = keras.layers.Input(shape=(4,))
o = keras.layers.dot([i1, i2], axes=1)
- self.assertListEqual(o.get_shape().as_list(), [None, 1])
+ self.assertListEqual(o.shape.as_list(), [None, 1])
model = keras.models.Model([i1, i2], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
_ = keras.layers.Dot(axes=1).get_config()
x1 = np.random.random((2, 4))
@@ -153,8 +139,9 @@
# Test with negative tuple of axes.
o = keras.layers.dot([i1, i2], axes=(-1, -1))
- self.assertListEqual(o.get_shape().as_list(), [None, 1])
+ self.assertListEqual(o.shape.as_list(), [None, 1])
model = keras.models.Model([i1, i2], o)
+ model.run_eagerly = testing_utils.should_run_eagerly()
out = model.predict([x1, x2])
self.assertEqual(out.shape, (2, 1))
self.assertAllClose(out, expected, atol=1e-4)
@@ -163,6 +150,32 @@
layer = keras.layers.Dot(axes=-1)
self.assertEqual(layer.compute_output_shape([(4, 5), (4, 5)]), (4, 1))
+
+@tf_test_util.run_all_in_graph_and_eager_modes
+class MergeLayersTestNoExecution(test.TestCase):
+
+ def test_merge_elementwise_errors(self):
+ i1 = keras.layers.Input(shape=(4, 5))
+ i2 = keras.layers.Input(shape=(4, 6))
+ with self.assertRaises(ValueError):
+ keras.layers.add([i1, i2])
+ with self.assertRaises(ValueError):
+ keras.layers.add([i1])
+ with self.assertRaises(ValueError):
+ keras.layers.add(i1)
+ with self.assertRaises(ValueError):
+ keras.layers.add([i1])
+
+ def test_concatenate_errors(self):
+ i1 = keras.layers.Input(shape=(4, 5))
+ i2 = keras.layers.Input(shape=(3, 5))
+ with self.assertRaisesRegexp(ValueError, 'inputs with matching shapes'):
+ keras.layers.concatenate([i1, i2], axis=-1)
+ with self.assertRaisesRegexp(ValueError, 'called on a list'):
+ keras.layers.concatenate(i1, axis=-1)
+ with self.assertRaisesRegexp(ValueError, 'called on a list'):
+ keras.layers.concatenate([i1], axis=-1)
+
def test_dot_errors(self):
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 6))
@@ -183,7 +196,7 @@
i1 = keras.layers.Input(shape=(4, 5))
i2 = keras.layers.Input(shape=(4, 5))
y = keras.layers.subtract([i1, i2])
- self.assertEqual(y.get_shape().as_list(), [None, 4, 5])
+ self.assertEqual(y.shape.as_list(), [None, 4, 5])
# Test invalid use cases
i1 = keras.layers.Input(shape=(4, 5))
@@ -193,39 +206,32 @@
with self.assertRaises(ValueError):
keras.layers.subtract([i1, i1, i1])
-
-class MergeLayersGraphOnlyTest(test.TestCase):
-
def test_merge_add_masking(self):
- with self.cached_session():
- i1 = keras.layers.Input(shape=(4, 5))
- i2 = keras.layers.Input(shape=(4, 5))
- m1 = keras.layers.Masking()(i1)
- layer = keras.layers.Add()
- o = layer([m1, i2])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 5])
- mask = layer.output_mask
- self.assertListEqual(mask.get_shape().as_list(), [None, 4])
+ i1 = keras.layers.Input(shape=(4, 5))
+ i2 = keras.layers.Input(shape=(4, 5))
+ m1 = keras.layers.Masking()(i1)
+ layer = keras.layers.Add()
+ o = layer([m1, i2])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 5])
+ mask = layer.output_mask
+ self.assertListEqual(mask.shape.as_list(), [None, 4])
- @tf_test_util.run_deprecated_v1
def test_merge_add_dynamic_shape(self):
- with self.cached_session():
- i1 = array_ops.placeholder(shape=(4, None), dtype='float32')
- i2 = array_ops.placeholder(shape=(4, 5), dtype='float32')
- layer = keras.layers.Add()
- o = layer([i1, i2])
- self.assertListEqual(o.get_shape().as_list(), [4, 5])
+ i1 = keras.Input(batch_shape=(4, None), dtype='float32')
+ i2 = keras.Input(batch_shape=(4, 5), dtype='float32')
+ layer = keras.layers.Add()
+ o = layer([i1, i2])
+ self.assertListEqual(o.shape.as_list(), [4, 5])
def test_merge_concatenate_masking(self):
- with self.cached_session():
- i1 = keras.layers.Input(shape=(4, 5))
- i2 = keras.layers.Input(shape=(4, 5))
- m1 = keras.layers.Masking()(i1)
- layer = keras.layers.Concatenate()
- o = layer([m1, i2])
- self.assertListEqual(o.get_shape().as_list(), [None, 4, 10])
- mask = layer.output_mask
- self.assertListEqual(mask.get_shape().as_list(), [None, 4])
+ i1 = keras.layers.Input(shape=(4, 5))
+ i2 = keras.layers.Input(shape=(4, 5))
+ m1 = keras.layers.Masking()(i1)
+ layer = keras.layers.Concatenate()
+ o = layer([m1, i2])
+ self.assertListEqual(o.shape.as_list(), [None, 4, 10])
+ mask = layer.output_mask
+ self.assertListEqual(mask.shape.as_list(), [None, 4])
if __name__ == '__main__':
diff --git a/tensorflow/python/keras/layers/noise_test.py b/tensorflow/python/keras/layers/noise_test.py
index 325dd93..f1537a6 100644
--- a/tensorflow/python/keras/layers/noise_test.py
+++ b/tensorflow/python/keras/layers/noise_test.py
@@ -19,13 +19,13 @@
from __future__ import print_function
from tensorflow.python import keras
-from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.platform import test
-@tf_test_util.run_all_in_graph_and_eager_modes
-class NoiseLayersTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class NoiseLayersTest(keras_parameterized.TestCase):
def test_GaussianNoise(self):
testing_utils.layer_test(
diff --git a/tensorflow/python/keras/layers/normalization.py b/tensorflow/python/keras/layers/normalization.py
index d958497..ee37e8a 100644
--- a/tensorflow/python/keras/layers/normalization.py
+++ b/tensorflow/python/keras/layers/normalization.py
@@ -18,7 +18,10 @@
from __future__ import division
from __future__ import print_function
+import contextlib
+
from tensorflow.python import tf2
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
@@ -37,7 +40,6 @@
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables as tf_variables
from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.util.tf_export import tf_export
@@ -412,11 +414,19 @@
def _assign_moving_average(self, variable, value, momentum):
with ops.name_scope(None, 'AssignMovingAvg',
[variable, value, momentum]) as scope:
- with ops.colocate_with(variable):
+ # TODO(apassos,srbs,skyewm): the colocation constraints here are disabled
+ # because of a bug which leads cond_v2 to skip rewriting them creating
+ # conflicts.
+ if tf2.enabled():
+ cm = contextlib.contextmanager(lambda: (yield))()
+ else:
+ cm = ops.colocate_with(variable)
+ with cm:
decay = ops.convert_to_tensor(1.0 - momentum, name='decay')
if decay.dtype != variable.dtype.base_dtype:
decay = math_ops.cast(decay, variable.dtype.base_dtype)
- update_delta = (variable - math_ops.cast(value, variable.dtype)) * decay
+ update_delta = (
+ variable - math_ops.cast(value, variable.dtype)) * decay
return state_ops.assign_sub(variable, update_delta, name=scope)
def _fused_batch_norm(self, inputs, training):
diff --git a/tensorflow/python/keras/layers/normalization_test.py b/tensorflow/python/keras/layers/normalization_test.py
index 9138c0a..780e02c 100644
--- a/tensorflow/python/keras/layers/normalization_test.py
+++ b/tensorflow/python/keras/layers/normalization_test.py
@@ -22,15 +22,16 @@
from tensorflow.python import keras
from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.keras.layers import normalization
from tensorflow.python.platform import test
from tensorflow.python.training import gradient_descent
-@tf_test_util.run_all_in_graph_and_eager_modes
-class NormalizationLayersTest(test.TestCase):
+class BatchNormalizationTest(keras_parameterized.TestCase):
+ @keras_parameterized.run_all_keras_modes
def test_basic_batchnorm(self):
testing_utils.layer_test(
keras.layers.BatchNormalization,
@@ -55,15 +56,8 @@
kwargs={'scale': False,
'center': False},
input_shape=(3, 3))
- testing_utils.layer_test(
- normalization.BatchNormalizationV2,
- kwargs={'fused': True},
- input_shape=(3, 3, 3, 3))
- testing_utils.layer_test(
- normalization.BatchNormalizationV2,
- kwargs={'fused': None},
- input_shape=(3, 3, 3))
+ @tf_test_util.run_in_graph_and_eager_modes
def test_batchnorm_weights(self):
layer = keras.layers.BatchNormalization(scale=False, center=False)
layer.build((None, 3, 4))
@@ -75,6 +69,7 @@
self.assertEqual(len(layer.trainable_weights), 2)
self.assertEqual(len(layer.weights), 4)
+ @tf_test_util.run_in_graph_and_eager_modes
def test_batchnorm_regularization(self):
layer = keras.layers.BatchNormalization(
gamma_regularizer='l1', beta_regularizer='l1')
@@ -87,36 +82,7 @@
self.assertEqual(layer.gamma.constraint, max_norm)
self.assertEqual(layer.beta.constraint, max_norm)
- def _test_batchnorm_correctness(self, dtype, use_v2=True, fused=False):
- model = keras.models.Sequential()
- layer_ctor = (normalization.BatchNormalizationV2 if use_v2
- else normalization.BatchNormalizationV1)
- norm = layer_ctor(input_shape=(2, 2, 2), momentum=0.8, fused=fused)
- model.add(norm)
- model.compile(loss='mse',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
-
- # centered on 5.0, variance 10.0
- x = (np.random.normal(loc=5.0, scale=10.0, size=(1000, 2, 2, 2))
- .astype(dtype))
- model.fit(x, x, epochs=4, verbose=0)
- out = model.predict(x)
- out -= keras.backend.eval(norm.beta)
- out /= keras.backend.eval(norm.gamma)
-
- np.testing.assert_allclose(out.mean(), 0.0, atol=1e-1)
- np.testing.assert_allclose(out.std(), 1.0, atol=1e-1)
-
- def test_batchnorm_correctness(self):
- self._test_batchnorm_correctness(np.float32)
- self._test_batchnorm_correctness(np.float32, fused=True)
- self._test_batchnorm_correctness(np.float32, use_v2=False)
-
- def test_batchnorm_mixed_precision(self):
- self._test_batchnorm_correctness(np.float16)
- self._test_batchnorm_correctness(np.float16, fused=True)
- self._test_batchnorm_correctness(np.float16, use_v2=False)
-
+ @keras_parameterized.run_all_keras_modes
def test_batchnorm_convnet(self):
if test.is_gpu_available(cuda_only=True):
with self.session(use_gpu=True):
@@ -125,7 +91,8 @@
axis=1, input_shape=(3, 4, 4), momentum=0.8)
model.add(norm)
model.compile(loss='mse',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01),
+ run_eagerly=testing_utils.should_run_eagerly())
# centered on 5.0, variance 10.0
x = np.random.normal(loc=5.0, scale=10.0, size=(1000, 3, 4, 4))
@@ -137,13 +104,15 @@
np.testing.assert_allclose(np.mean(out, axis=(0, 2, 3)), 0.0, atol=1e-1)
np.testing.assert_allclose(np.std(out, axis=(0, 2, 3)), 1.0, atol=1e-1)
+ @keras_parameterized.run_all_keras_modes
def test_batchnorm_convnet_channel_last(self):
model = keras.models.Sequential()
norm = keras.layers.BatchNormalization(
axis=-1, input_shape=(4, 4, 3), momentum=0.8)
model.add(norm)
model.compile(loss='mse',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01),
+ run_eagerly=testing_utils.should_run_eagerly())
# centered on 5.0, variance 10.0
x = np.random.normal(loc=5.0, scale=10.0, size=(1000, 4, 4, 3))
@@ -155,6 +124,28 @@
np.testing.assert_allclose(np.mean(out, axis=(0, 1, 2)), 0.0, atol=1e-1)
np.testing.assert_allclose(np.std(out, axis=(0, 1, 2)), 1.0, atol=1e-1)
+ @keras_parameterized.run_all_keras_modes
+ def test_batchnorm_correctness(self):
+ _run_batchnorm_correctness_test(
+ normalization.BatchNormalization, dtype='float32')
+ _run_batchnorm_correctness_test(
+ normalization.BatchNormalization, dtype='float32', fused=True)
+ _run_batchnorm_correctness_test(
+ normalization.BatchNormalization, dtype='float32', fused=False)
+
+ @keras_parameterized.run_all_keras_modes
+ def test_batchnorm_mixed_precision(self):
+ _run_batchnorm_correctness_test(
+ normalization.BatchNormalization, dtype='float16')
+ _run_batchnorm_correctness_test(
+ normalization.BatchNormalization, dtype='float16', fused=True)
+ _run_batchnorm_correctness_test(
+ normalization.BatchNormalization, dtype='float16', fused=False)
+
+
+class BatchNormalizationV1Test(test.TestCase):
+
+ @tf_test_util.run_in_graph_and_eager_modes
def test_v1_fused_attribute(self):
norm = normalization.BatchNormalizationV1()
inp = keras.layers.Input((4, 4, 4))
@@ -173,6 +164,21 @@
norm(inp)
self.assertEqual(norm.fused, False)
+
+class BatchNormalizationV2Test(keras_parameterized.TestCase):
+
+ @keras_parameterized.run_all_keras_modes
+ def test_basic_batchnorm_v2(self):
+ testing_utils.layer_test(
+ normalization.BatchNormalizationV2,
+ kwargs={'fused': True},
+ input_shape=(3, 3, 3, 3))
+ testing_utils.layer_test(
+ normalization.BatchNormalizationV2,
+ kwargs={'fused': None},
+ input_shape=(3, 3, 3))
+
+ @tf_test_util.run_in_graph_and_eager_modes
def test_v2_fused_attribute(self):
norm = normalization.BatchNormalizationV2()
self.assertEqual(norm.fused, None)
@@ -227,6 +233,27 @@
norm(inp)
+def _run_batchnorm_correctness_test(layer, dtype='float32', fused=False):
+ model = keras.models.Sequential()
+ norm = layer(input_shape=(2, 2, 2), momentum=0.8, fused=fused)
+ model.add(norm)
+ model.compile(loss='mse',
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01),
+ run_eagerly=testing_utils.should_run_eagerly())
+
+ # centered on 5.0, variance 10.0
+ x = (np.random.normal(loc=5.0, scale=10.0, size=(1000, 2, 2, 2))
+ .astype(dtype))
+ model.fit(x, x, epochs=4, verbose=0)
+ out = model.predict(x)
+ out -= keras.backend.eval(norm.beta)
+ out /= keras.backend.eval(norm.gamma)
+
+ np.testing.assert_allclose(out.mean(), 0.0, atol=1e-1)
+ np.testing.assert_allclose(out.std(), 1.0, atol=1e-1)
+
+
+@tf_test_util.run_v1_only('b/120545219')
class NormalizationLayersGraphModeOnlyTest(test.TestCase):
def test_shared_batchnorm(self):
@@ -301,13 +328,14 @@
x2 = model.predict(val_a)
self.assertAllClose(x1, x2, atol=1e-7)
- @tf_test_util.run_deprecated_v1
def test_batchnorm_trainable(self):
"""Tests that batchnorm layer is trainable when learning phase is enabled.
Computes mean and std for current inputs then
applies batch normalization using them.
"""
+ # TODO(fchollet): enable in all execution modes when issue with
+ # learning phase setting is resolved.
with self.cached_session():
bn_mean = 0.5
bn_std = 10.
diff --git a/tensorflow/python/keras/layers/recurrent.py b/tensorflow/python/keras/layers/recurrent.py
index 189ad98..568e879 100644
--- a/tensorflow/python/keras/layers/recurrent.py
+++ b/tensorflow/python/keras/layers/recurrent.py
@@ -1497,12 +1497,6 @@
initializer=self.bias_initializer,
regularizer=self.bias_regularizer,
constraint=self.bias_constraint)
- if not self.reset_after:
- self.input_bias, self.recurrent_bias = self.bias, None
- else:
- self.input_bias = K.flatten(self.bias[0])
- self.recurrent_bias = K.flatten(self.bias[1])
-
else:
self.bias = None
self.built = True
@@ -1529,6 +1523,12 @@
# dropout matrices for recurrent units
rec_dp_mask = self._recurrent_dropout_mask
+ if self.use_bias:
+ if not self.reset_after:
+ input_bias, recurrent_bias = self.bias, None
+ else:
+ input_bias, recurrent_bias = array_ops.unstack(self.bias)
+
if self.implementation == 1:
if 0. < self.dropout < 1.:
inputs_z = inputs * dp_mask[0]
@@ -1544,9 +1544,9 @@
x_h = K.dot(inputs_h, self.kernel[:, self.units * 2:])
if self.use_bias:
- x_z = K.bias_add(x_z, self.input_bias[:self.units])
- x_r = K.bias_add(x_r, self.input_bias[self.units: self.units * 2])
- x_h = K.bias_add(x_h, self.input_bias[self.units * 2:])
+ x_z = K.bias_add(x_z, input_bias[:self.units])
+ x_r = K.bias_add(x_r, input_bias[self.units: self.units * 2])
+ x_h = K.bias_add(x_h, input_bias[self.units * 2:])
if 0. < self.recurrent_dropout < 1.:
h_tm1_z = h_tm1 * rec_dp_mask[0]
@@ -1561,10 +1561,9 @@
recurrent_r = K.dot(h_tm1_r,
self.recurrent_kernel[:, self.units:self.units * 2])
if self.reset_after and self.use_bias:
- recurrent_z = K.bias_add(recurrent_z, self.recurrent_bias[:self.units])
+ recurrent_z = K.bias_add(recurrent_z, recurrent_bias[:self.units])
recurrent_r = K.bias_add(recurrent_r,
- self.recurrent_bias[self.units:
- self.units * 2])
+ recurrent_bias[self.units:self.units * 2])
z = self.recurrent_activation(x_z + recurrent_z)
r = self.recurrent_activation(x_r + recurrent_r)
@@ -1573,8 +1572,7 @@
if self.reset_after:
recurrent_h = K.dot(h_tm1_h, self.recurrent_kernel[:, self.units * 2:])
if self.use_bias:
- recurrent_h = K.bias_add(recurrent_h,
- self.recurrent_bias[self.units * 2:])
+ recurrent_h = K.bias_add(recurrent_h, recurrent_bias[self.units * 2:])
recurrent_h = r * recurrent_h
else:
recurrent_h = K.dot(r * h_tm1_h,
@@ -1589,7 +1587,7 @@
matrix_x = K.dot(inputs, self.kernel)
if self.use_bias:
# biases: bias_z_i, bias_r_i, bias_h_i
- matrix_x = K.bias_add(matrix_x, self.input_bias)
+ matrix_x = K.bias_add(matrix_x, input_bias)
x_z = matrix_x[:, :self.units]
x_r = matrix_x[:, self.units: 2 * self.units]
@@ -1602,7 +1600,7 @@
# hidden state projected by all gate matrices at once
matrix_inner = K.dot(h_tm1, self.recurrent_kernel)
if self.use_bias:
- matrix_inner = K.bias_add(matrix_inner, self.recurrent_bias)
+ matrix_inner = K.bias_add(matrix_inner, recurrent_bias)
else:
# hidden state projected separately for update/reset and new
matrix_inner = K.dot(h_tm1, self.recurrent_kernel[:, :2 * self.units])
@@ -2274,7 +2272,7 @@
return c, o
-@tf_export('keras.layers.LSTM')
+@tf_export(v1=['keras.layers.LSTM'])
class LSTM(RNN):
"""Long Short-Term Memory layer - Hochreiter 1997.
@@ -2533,6 +2531,7 @@
return cls(**config)
+@tf_export('keras.layers.LSTM', v1=[])
class UnifiedLSTM(LSTM):
"""Long Short-Term Memory layer - Hochreiter 1997.
@@ -2546,13 +2545,11 @@
Arguments:
units: Positive integer, dimensionality of the output space.
activation: Activation function to use.
- Default: hyperbolic tangent (`tanh`). If you pass `None`, no activation
- is applied
- (ie. "linear" activation: `a(x) = x`).
+ Default: hyperbolic tangent (`tanh`). If you pass `None`, no activation
+ is applied (ie. "linear" activation: `a(x) = x`).
recurrent_activation: Activation function to use for the recurrent step.
- Default: hard sigmoid (`hard_sigmoid`). If you pass `None`, no
- activation is applied
- (ie. "linear" activation: `a(x) = x`).
+ Default: sigmoid (`sigmoid`). If you pass `None`, no activation is
+ applied (ie. "linear" activation: `a(x) = x`).
use_bias: Boolean, whether the layer uses a bias vector.
kernel_initializer: Initializer for the `kernel` weights matrix, used for
the linear transformation of the inputs..
@@ -2602,7 +2599,7 @@
def __init__(self,
units,
activation='tanh',
- recurrent_activation='hard_sigmoid',
+ recurrent_activation='sigmoid',
use_bias=True,
kernel_initializer='glorot_uniform',
recurrent_initializer='orthogonal',
@@ -2659,27 +2656,10 @@
self.state_spec = [
InputSpec(shape=(None, dim)) for dim in (self.units, self.units)
]
- self._num_constants = None
- self._num_inputs = None
self._dropout_mask = None
self.could_use_cudnn = (
- activation == 'tanh' and recurrent_dropout == 0 and
- not unroll and use_bias and bias_regularizer is None)
-
- def build(self, input_shape):
- super(UnifiedLSTM, self).build(input_shape)
- if self.could_use_cudnn:
- # Add a new set of bias for CuDNN implementation only. Standard LSTM only
- # has bias for recurrent kernel, while CuDNN LSTM has an extra set for
- # input gate as well.
- self.cudnn_bias = self.add_weight(
- shape=(self.units * 4,),
- name='cudnn_bias',
- use_resource=True,
- initializer=self.bias_initializer,
- regularizer=self.bias_regularizer,
- constraint=self.bias_constraint)
- self.built = True
+ activation == 'tanh' and recurrent_activation == 'sigmoid' and
+ recurrent_dropout == 0 and not unroll and use_bias)
def call(self, inputs, mask=None, training=None, initial_state=None):
# LSTM does not support constants. Ignore it during process.
@@ -2719,9 +2699,7 @@
# both normal and CuDNN implementations.
if self.go_backwards:
# Reverse time axis.
- inputs = K.reverse(inputs, 1)
-
- combined_bias = array_ops.concat([self.cudnn_bias, self.cell.bias], 0)
+ inputs = K.reverse(inputs, 0 if self.time_major else 1)
if 0 < self.dropout < 1:
if self._dropout_mask is None:
@@ -2757,23 +2735,23 @@
if context.num_gpus() > 0:
last_output, outputs, new_h, new_c, runtime = defun_cudnn_lstm(
inputs, initial_state[0], initial_state[1], self.cell.kernel,
- self.cell.recurrent_kernel, combined_bias, self.time_major)
+ self.cell.recurrent_kernel, self.cell.bias, self.time_major)
else:
last_output, outputs, new_h, new_c, runtime = defun_standard_lstm(
inputs, initial_state[0], initial_state[1], self.cell.kernel,
- self.cell.recurrent_kernel, combined_bias, self.activation,
+ self.cell.recurrent_kernel, self.cell.bias, self.activation,
self.recurrent_activation, self.time_major)
else:
# Call the normal LSTM impl and register the CuDNN impl function. The
# grappler will kick in during session execution to optimize the graph.
last_output, outputs, new_h, new_c, runtime = defun_standard_lstm(
inputs, initial_state[0], initial_state[1], self.cell.kernel,
- self.cell.recurrent_kernel, combined_bias, self.activation,
+ self.cell.recurrent_kernel, self.cell.bias, self.activation,
self.recurrent_activation, self.time_major)
function.register(defun_cudnn_lstm, inputs, initial_state[0],
initial_state[1], self.cell.kernel,
- self.cell.recurrent_kernel, combined_bias,
+ self.cell.recurrent_kernel, self.cell.bias,
self.time_major)
states = [new_h, new_c]
@@ -2795,56 +2773,6 @@
else:
return output
- @property
- def trainable_weights(self):
- if self.trainable:
- weights = []
- weights += self.cell.trainable_weights
- if getattr(self, 'cudnn_bias', None) is not None:
- weights += [self.cudnn_bias]
- return weights
- return []
-
- @property
- def non_trainable_weights(self):
- if not self.trainable:
- weights = []
- weights += self.cell.non_trainable_weights
- if getattr(self, 'cudnn_bias', None) is not None:
- weights += [self.cudnn_bias]
- return weights
- return []
-
- @property
- def losses(self):
- losses = []
- losses += self.cell.losses
- return losses + self._losses
-
- @property
- def updates(self):
- updates = []
- updates += self.cell.updates
- return updates + self._updates
-
- def get_weights(self):
- weights = []
- weights += self.cell.weights
- if getattr(self, 'cudnn_bias', None) is not None:
- weights += [self.cudnn_bias]
- return K.batch_get_value(weights)
-
- def set_weights(self, weights):
- tuples = []
- cell_weights = weights[:len(self.cell.weights)]
- if cell_weights:
- tuples.append((self.cell.weights, cell_weights))
- if getattr(self, 'cudnn_bias', None) is not None:
- cudnn_bias_weights = weights[len(self.cell.weights):]
- if cudnn_bias_weights:
- tuples.append((self.cudnn_bias, cudnn_bias_weights))
- K.batch_set_value(tuples)
-
def _canonical_to_params(weights, biases, shape, transpose_weights=False):
"""Utility function convert variable to CuDNN compatible parameter.
@@ -2919,9 +2847,6 @@
input_shape = K.int_shape(inputs)
timesteps = input_shape[0] if time_major else input_shape[1]
- # Only use the second half of the bias weights.
- _, real_bias = array_ops.split(bias, 2)
-
def step(cell_inputs, cell_states):
"""Step function that will be used by Keras RNN backend."""
h_tm1 = cell_states[0] # previous memory state
@@ -2929,7 +2854,7 @@
z = K.dot(cell_inputs, kernel)
z += K.dot(h_tm1, recurrent_kernel)
- z = K.bias_add(z, real_bias)
+ z = K.bias_add(z, bias)
z0, z1, z2, z3 = array_ops.split(z, 4, axis=1)
@@ -2962,9 +2887,13 @@
weights = array_ops.split(kernel, 4, axis=1)
weights += array_ops.split(recurrent_kernel, 4, axis=1)
+ # CuDNN has an extra set of bias for inputs, we disable them (setting to 0),
+ # so that mathematically it is same as the canonical LSTM implementation.
+ full_bias = array_ops.concat((array_ops.zeros_like(bias), bias), 0)
+
params = _canonical_to_params(
weights=weights,
- biases=array_ops.split(bias, 8),
+ biases=array_ops.split(full_bias, 8),
shape=constant_op.constant([-1]),
transpose_weights=True)
diff --git a/tensorflow/python/keras/layers/recurrent_test.py b/tensorflow/python/keras/layers/recurrent_test.py
index b144906..e9bf788 100644
--- a/tensorflow/python/keras/layers/recurrent_test.py
+++ b/tensorflow/python/keras/layers/recurrent_test.py
@@ -30,7 +30,8 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import tensor_shape
-from tensorflow.python.framework import test_util
+from tensorflow.python.keras import keras_parameterized
+from tensorflow.python.keras import testing_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
@@ -48,8 +49,8 @@
NestedState = collections.namedtuple('NestedState', ['s1', 's2'])
-@test_util.run_all_in_graph_and_eager_modes
-class RNNTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class RNNTest(keras_parameterized.TestCase):
def test_minimal_rnn_cell_non_layer(self):
@@ -73,7 +74,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))
# Test stacking.
@@ -84,7 +86,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))
def test_minimal_rnn_cell_non_layer_multiple_states(self):
@@ -112,7 +115,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))
# Test stacking.
@@ -125,7 +129,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))
def test_minimal_rnn_cell_layer(self):
@@ -165,7 +170,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))
# Test basic case serialization.
@@ -189,7 +195,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))
# Test stacked RNN serialization.
@@ -228,7 +235,8 @@
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, embedding_dim)),
np.zeros((batch, time_step, units)))
@@ -246,7 +254,8 @@
y = keras.layers.Lambda(lambda t: array_ops.transpose(t, [1, 0, 2]))(y)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, embedding_dim)),
np.zeros((batch, time_step, cell_units[-1])))
@@ -261,7 +270,8 @@
y = keras.layers.Lambda(lambda t: array_ops.transpose(t, [1, 0, 2]))(rnn)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, embedding_dim)),
np.zeros((batch, time_step, units)))
@@ -273,7 +283,8 @@
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, embedding_dim)),
np.zeros((batch, time_step, units)))
@@ -347,7 +358,8 @@
model = keras.models.Model([x, c], y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((6, 5, 5)), np.zeros((6, 3))],
np.zeros((6, 32))
@@ -385,7 +397,8 @@
y = layer(x, constants=c)
model = keras.models.Model([x, c], y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((6, 5, 5)), np.zeros((6, 3))],
np.zeros((6, 32))
@@ -399,7 +412,8 @@
y = layer(x, constants=c)
model = keras.models.Model([x, c], y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((6, 5, 5)), np.zeros((6, 3))],
np.zeros((6, 32))
@@ -471,7 +485,8 @@
y = layer(x, initial_state=s, constants=c)
model = keras.models.Model([x, s, c], y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((6, 5, 5)), np.zeros((6, 32)), np.zeros((6, 3))],
np.zeros((6, 32))
@@ -601,7 +616,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
# Test basic case serialization.
x_np = np.random.random((6, 5, 5))
@@ -623,7 +639,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
# Test stacked RNN serialization.
x_np = np.random.random((6, 5, 5))
@@ -647,7 +664,7 @@
x = keras.Input((None, 5))
y = layer(x)
model = keras.models.Model(x, y)
- model.compile('sgd', 'mse')
+ model.compile('sgd', 'mse', run_eagerly=testing_utils.should_run_eagerly())
x_np = np.random.random((6, 5, 5))
y_np = np.random.random((6, 3))
model.train_on_batch(x_np, y_np)
@@ -690,7 +707,8 @@
model = keras.models.Sequential()
model.add(rnn(2))
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.fit(x, y, epochs=1, batch_size=1)
# check whether the model variables are present in the
@@ -723,7 +741,8 @@
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, input_a, input_b)),
np.zeros((batch, unit_a, unit_b)))
@@ -739,7 +758,8 @@
y = layer(x)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, input_a, input_b)),
np.zeros((batch, unit_a * 4, unit_b * 4)))
@@ -762,7 +782,8 @@
model = keras.models.Model([x, s], y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch([
np.zeros((batch, time_step, input_a, input_b)),
np.zeros((batch, unit_a, unit_b))
@@ -799,7 +820,8 @@
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
np.zeros((batch, time_step, input_size)),
np.zeros((batch, input_size)))
@@ -854,7 +876,8 @@
model = keras.models.Model((input_1, input_2), outputs)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((batch, t, i1)), np.zeros((batch, t, i2, i3))],
[np.zeros((batch, o1)), np.zeros((batch, o2, o3))])
@@ -875,7 +898,8 @@
model = keras.models.Model([input_1, input_2], outputs)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((batch, t, i1)),
np.zeros((batch, t, i2, i3))],
@@ -903,7 +927,8 @@
model = keras.models.Model([input_1, input_2], [output1, output2])
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((batch, t, i1)),
np.zeros((batch, t, i2, i3))],
@@ -927,7 +952,8 @@
model = keras.models.Model([input_1, input_2], [output1, output2])
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((batch, t, i1)),
np.zeros((batch, t, i2, i3))],
@@ -960,7 +986,8 @@
model = keras.models.Model([input_1, input_2, init_s1, init_s2],
[output1, output2])
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((batch, t, i1)),
np.zeros((batch, t, i2, i3)),
@@ -991,7 +1018,8 @@
model = keras.models.Model([input_1, input_2, init_s1, init_s2],
[output1, output2])
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
model.train_on_batch(
[np.zeros((batch, t, i1)),
np.zeros((batch, t, i2, i3)),
@@ -1004,18 +1032,17 @@
def test_peephole_lstm_cell(self):
def _run_cell(cell_fn, **kwargs):
- with self.cached_session() as sess:
- inputs = array_ops.one_hot([1, 2, 3, 4], 4)
- cell = cell_fn(5, **kwargs)
- cell.build(inputs.shape)
- initial_state = cell.get_initial_state(
- inputs=inputs, batch_size=4, dtype=dtypes.float32)
- inputs, _ = cell(inputs, initial_state)
- output = inputs
- if not context.executing_eagerly():
- self.evaluate(variables_lib.global_variables_initializer())
- output = self.evaluate(output)
- return output
+ inputs = array_ops.one_hot([1, 2, 3, 4], 4)
+ cell = cell_fn(5, **kwargs)
+ cell.build(inputs.shape)
+ initial_state = cell.get_initial_state(
+ inputs=inputs, batch_size=4, dtype=dtypes.float32)
+ inputs, _ = cell(inputs, initial_state)
+ output = inputs
+ if not context.executing_eagerly():
+ self.evaluate(variables_lib.global_variables_initializer())
+ output = self.evaluate(output)
+ return output
random_seed.set_random_seed(12345)
# `recurrent_activation` kwarg is set to sigmoid as that is hardcoded into
@@ -1067,7 +1094,8 @@
Cell(), return_state=True)(x_masked, initial_state=s_0)
model = keras.models.Model([x, s_0], [y, s])
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
# last time step masked
x_np = np.array([[[1.], [2.], [0.]]])
@@ -1091,7 +1119,8 @@
y = layer(masked_input)
model = keras.models.Model(x, y)
model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001),
- loss='mse')
+ loss='mse',
+ run_eagerly=testing_utils.should_run_eagerly())
np_x = np.ones((6, 5, 5))
result_1 = model.predict(np_x)
diff --git a/tensorflow/python/keras/layers/simplernn_test.py b/tensorflow/python/keras/layers/simplernn_test.py
index b49b159..58f2f9a 100644
--- a/tensorflow/python/keras/layers/simplernn_test.py
+++ b/tensorflow/python/keras/layers/simplernn_test.py
@@ -22,14 +22,15 @@
from tensorflow.python import keras
from tensorflow.python.framework import test_util as tf_test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.platform import test
from tensorflow.python.training import gradient_descent
from tensorflow.python.training.rmsprop import RMSPropOptimizer
-@tf_test_util.run_all_in_graph_and_eager_modes
-class SimpleRNNLayerTest(test.TestCase):
+@keras_parameterized.run_all_keras_modes
+class SimpleRNNLayerTest(keras_parameterized.TestCase):
def test_return_sequences_SimpleRNN(self):
num_samples = 2
@@ -98,7 +99,7 @@
self.assertEqual(layer.cell.recurrent_kernel.constraint, r_constraint)
self.assertEqual(layer.cell.bias.constraint, b_constraint)
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_with_masking_layer_SimpleRNN(self):
layer_class = keras.layers.SimpleRNN
inputs = np.random.random((2, 3, 4))
@@ -118,93 +119,91 @@
l2 = layer_class.from_config(l1.get_config())
assert l1.get_config() == l2.get_config()
-
-class SimpleRNNLayerGraphOnlyTest(test.TestCase):
-
- @tf_test_util.run_deprecated_v1
def test_statefulness_SimpleRNN(self):
num_samples = 2
timesteps = 3
embedding_dim = 4
units = 2
layer_class = keras.layers.SimpleRNN
- with self.cached_session():
- model = keras.models.Sequential()
- model.add(
- keras.layers.Embedding(
- 4,
- embedding_dim,
- mask_zero=True,
- input_length=timesteps,
- batch_input_shape=(num_samples, timesteps)))
- layer = layer_class(
- units, return_sequences=False, stateful=True, weights=None)
- model.add(layer)
- model.compile(optimizer=gradient_descent.GradientDescentOptimizer(0.01),
- loss='mse')
- out1 = model.predict(np.ones((num_samples, timesteps)))
- self.assertEqual(out1.shape, (num_samples, units))
+ model = keras.models.Sequential()
+ model.add(
+ keras.layers.Embedding(
+ 4,
+ embedding_dim,
+ mask_zero=True,
+ input_length=timesteps,
+ batch_input_shape=(num_samples, timesteps)))
+ layer = layer_class(
+ units, return_sequences=False, stateful=True, weights=None)
+ model.add(layer)
+ model.compile(optimizer=gradient_descent.GradientDescentOptimizer(0.01),
+ loss='mse')
+ out1 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertEqual(out1.shape, (num_samples, units))
- # train once so that the states change
- model.train_on_batch(
- np.ones((num_samples, timesteps)), np.ones((num_samples, units)))
- out2 = model.predict(np.ones((num_samples, timesteps)))
+ # train once so that the states change
+ model.train_on_batch(
+ np.ones((num_samples, timesteps)), np.ones((num_samples, units)))
+ out2 = model.predict(np.ones((num_samples, timesteps)))
- # if the state is not reset, output should be different
- self.assertNotEqual(out1.max(), out2.max())
+ # if the state is not reset, output should be different
+ self.assertNotEqual(out1.max(), out2.max())
- # check that output changes after states are reset
- # (even though the model itself didn't change)
- layer.reset_states()
- out3 = model.predict(np.ones((num_samples, timesteps)))
- self.assertNotEqual(out2.max(), out3.max())
+ # check that output changes after states are reset
+ # (even though the model itself didn't change)
+ layer.reset_states()
+ out3 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertNotEqual(out2.max(), out3.max())
- # check that container-level reset_states() works
- model.reset_states()
- out4 = model.predict(np.ones((num_samples, timesteps)))
- np.testing.assert_allclose(out3, out4, atol=1e-5)
+ # check that container-level reset_states() works
+ model.reset_states()
+ out4 = model.predict(np.ones((num_samples, timesteps)))
+ np.testing.assert_allclose(out3, out4, atol=1e-5)
- # check that the call to `predict` updated the states
- out5 = model.predict(np.ones((num_samples, timesteps)))
- self.assertNotEqual(out4.max(), out5.max())
+ # check that the call to `predict` updated the states
+ out5 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertNotEqual(out4.max(), out5.max())
- # Check masking
- layer.reset_states()
+ # Check masking
+ layer.reset_states()
- left_padded_input = np.ones((num_samples, timesteps))
- left_padded_input[0, :1] = 0
- left_padded_input[1, :2] = 0
- out6 = model.predict(left_padded_input)
+ left_padded_input = np.ones((num_samples, timesteps))
+ left_padded_input[0, :1] = 0
+ left_padded_input[1, :2] = 0
+ out6 = model.predict(left_padded_input)
- layer.reset_states()
+ layer.reset_states()
- right_padded_input = np.ones((num_samples, timesteps))
- right_padded_input[0, -1:] = 0
- right_padded_input[1, -2:] = 0
- out7 = model.predict(right_padded_input)
+ right_padded_input = np.ones((num_samples, timesteps))
+ right_padded_input[0, -1:] = 0
+ right_padded_input[1, -2:] = 0
+ out7 = model.predict(right_padded_input)
- np.testing.assert_allclose(out7, out6, atol=1e-5)
+ np.testing.assert_allclose(out7, out6, atol=1e-5)
+
+class SimpleRNNLayerGraphOnlyTest(test.TestCase):
+
+ # b/120919032
@tf_test_util.run_deprecated_v1
def test_regularizers_SimpleRNN(self):
embedding_dim = 4
layer_class = keras.layers.SimpleRNN
- with self.cached_session():
- layer = layer_class(
- 5,
- return_sequences=False,
- weights=None,
- input_shape=(None, embedding_dim),
- kernel_regularizer=keras.regularizers.l1(0.01),
- recurrent_regularizer=keras.regularizers.l1(0.01),
- bias_regularizer='l2',
- activity_regularizer='l1')
- layer.build((None, None, 2))
- self.assertEqual(len(layer.losses), 3)
+ layer = layer_class(
+ 5,
+ return_sequences=False,
+ weights=None,
+ input_shape=(None, embedding_dim),
+ kernel_regularizer=keras.regularizers.l1(0.01),
+ recurrent_regularizer=keras.regularizers.l1(0.01),
+ bias_regularizer='l2',
+ activity_regularizer='l1')
+ layer.build((None, None, 2))
+ self.assertEqual(len(layer.losses), 3)
- x = keras.backend.variable(np.ones((2, 3, 2)))
- layer(x)
- self.assertEqual(len(layer.get_losses_for(x)), 1)
+ x = keras.backend.variable(np.ones((2, 3, 2)))
+ layer(x)
+ self.assertEqual(len(layer.get_losses_for(x)), 1)
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/layers/unified_lstm_test.py b/tensorflow/python/keras/layers/unified_lstm_test.py
index b004284..c513046 100644
--- a/tensorflow/python/keras/layers/unified_lstm_test.py
+++ b/tensorflow/python/keras/layers/unified_lstm_test.py
@@ -18,6 +18,8 @@
from __future__ import division
from __future__ import print_function
+import os
+import shutil
import time
from absl.testing import parameterized
@@ -26,12 +28,13 @@
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python import keras
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
-from tensorflow.python.keras.layers.cudnn_recurrent import CuDNNLSTM
-from tensorflow.python.keras.layers.recurrent import UnifiedLSTM
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_math_ops
@@ -52,8 +55,251 @@
_config = config_pb2.ConfigProto(graph_options=_graph_options)
-class UnifiedLSTMTest(test.TestCase, parameterized.TestCase):
+@keras_parameterized.run_all_keras_modes(config=_config)
+class UnifiedLSTMTest(keras_parameterized.TestCase):
+ @parameterized.named_parameters(
+ ('non_tan_activation', 'relu', 'sigmoid', 0, False, True),
+ ('non_sigmoid_recur_activation', 'tanh', 'relu', 0, False, True),
+ ('use_recurrent_dropout', 'tanh', 'sigmoid', 0.1, False, True),
+ ('unroll', 'tanh', 'sigmoid', 0, True, True),
+ ('not_use_bias', 'tanh', 'sigmoid', 0, False, False),
+ )
+ def test_could_use_defun_backend(self, activation, recurrent_activation,
+ recurrent_dropout, unroll, use_bias):
+ layer = keras.layers.UnifiedLSTM(
+ 1,
+ activation=activation,
+ recurrent_activation=recurrent_activation,
+ recurrent_dropout=recurrent_dropout,
+ unroll=unroll,
+ use_bias=use_bias)
+ self.assertFalse(layer.could_use_cudnn)
+
+ def test_static_shape_inference_LSTM(self):
+ # Github issue: 15165
+ timesteps = 3
+ embedding_dim = 4
+ units = 2
+
+ model = keras.models.Sequential()
+ inputs = keras.layers.Dense(
+ embedding_dim, input_shape=(timesteps, embedding_dim))
+ model.add(inputs)
+ layer = keras.layers.UnifiedLSTM(units, return_sequences=True)
+ model.add(layer)
+ outputs = model.layers[-1].output
+ self.assertEqual(outputs.get_shape().as_list(), [None, timesteps, units])
+
+ def test_dynamic_behavior_LSTM(self):
+ num_samples = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 2
+ layer = keras.layers.UnifiedLSTM(units, input_shape=(None, embedding_dim))
+ model = keras.models.Sequential()
+ model.add(layer)
+ model.compile(gradient_descent.GradientDescentOptimizer(0.001), 'mse')
+ x = np.random.random((num_samples, timesteps, embedding_dim))
+ y = np.random.random((num_samples, units))
+ model.train_on_batch(x, y)
+
+ def test_stacking_LSTM(self):
+ inputs = np.random.random((2, 3, 4))
+ targets = np.abs(np.random.random((2, 3, 5)))
+ targets /= targets.sum(axis=-1, keepdims=True)
+ model = keras.models.Sequential()
+ model.add(keras.layers.UnifiedLSTM(10, return_sequences=True, unroll=False))
+ model.add(keras.layers.UnifiedLSTM(5, return_sequences=True, unroll=False))
+ model.compile(
+ loss='categorical_crossentropy',
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+ model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
+
+ def test_from_config_LSTM(self):
+ layer_class = keras.layers.UnifiedLSTM
+ for stateful in (False, True):
+ l1 = layer_class(units=1, stateful=stateful)
+ l2 = layer_class.from_config(l1.get_config())
+ assert l1.get_config() == l2.get_config()
+
+ def test_specify_initial_state_keras_tensor(self):
+ num_states = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+
+ # Test with Keras tensor
+ inputs = keras.Input((timesteps, embedding_dim))
+ initial_state = [keras.Input((units,)) for _ in range(num_states)]
+ layer = keras.layers.UnifiedLSTM(units)
+ if len(initial_state) == 1:
+ output = layer(inputs, initial_state=initial_state[0])
+ else:
+ output = layer(inputs, initial_state=initial_state)
+ assert initial_state[0] in layer._inbound_nodes[0].input_tensors
+
+ model = keras.models.Model([inputs] + initial_state, output)
+ model.compile(
+ loss='categorical_crossentropy',
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+
+ inputs = np.random.random((num_samples, timesteps, embedding_dim))
+ initial_state = [
+ np.random.random((num_samples, units)) for _ in range(num_states)
+ ]
+ targets = np.random.random((num_samples, units))
+ model.train_on_batch([inputs] + initial_state, targets)
+
+ def DISABLED_test_specify_initial_state_non_keras_tensor(self):
+ num_states = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+
+ # Test with non-Keras tensor
+ inputs = keras.Input((timesteps, embedding_dim))
+ initial_state = [
+ keras.backend.random_normal_variable((num_samples, units), 0, 1)
+ for _ in range(num_states)
+ ]
+ layer = keras.layers.UnifiedLSTM(units)
+ output = layer(inputs, initial_state=initial_state)
+
+ model = keras.models.Model(inputs, output)
+ model.compile(
+ loss='categorical_crossentropy',
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+
+ inputs = np.random.random((num_samples, timesteps, embedding_dim))
+ targets = np.random.random((num_samples, units))
+ model.train_on_batch(inputs, targets)
+
+ def test_reset_states_with_values(self):
+ num_states = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+
+ layer = keras.layers.UnifiedLSTM(units, stateful=True)
+ layer.build((num_samples, timesteps, embedding_dim))
+ layer.reset_states()
+ assert len(layer.states) == num_states
+ assert layer.states[0] is not None
+ self.assertAllClose(
+ keras.backend.eval(layer.states[0]),
+ np.zeros(keras.backend.int_shape(layer.states[0])),
+ atol=1e-4)
+ state_shapes = [keras.backend.int_shape(state) for state in layer.states]
+ values = [np.ones(shape) for shape in state_shapes]
+ if len(values) == 1:
+ values = values[0]
+ layer.reset_states(values)
+ self.assertAllClose(
+ keras.backend.eval(layer.states[0]),
+ np.ones(keras.backend.int_shape(layer.states[0])),
+ atol=1e-4)
+
+ # Test with invalid data
+ with self.assertRaises(ValueError):
+ layer.reset_states([1] * (len(layer.states) + 1))
+
+ def test_specify_state_with_masking(self):
+ num_states = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+
+ inputs = keras.Input((timesteps, embedding_dim))
+ _ = keras.layers.Masking()(inputs)
+ initial_state = [keras.Input((units,)) for _ in range(num_states)]
+ output = keras.layers.UnifiedLSTM(units)(
+ inputs, initial_state=initial_state)
+
+ model = keras.models.Model([inputs] + initial_state, output)
+ model.compile(
+ loss='categorical_crossentropy',
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+
+ inputs = np.random.random((num_samples, timesteps, embedding_dim))
+ initial_state = [
+ np.random.random((num_samples, units)) for _ in range(num_states)
+ ]
+ targets = np.random.random((num_samples, units))
+ model.train_on_batch([inputs] + initial_state, targets)
+
+ def test_return_state(self):
+ num_states = 2
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+
+ inputs = keras.Input(batch_shape=(num_samples, timesteps, embedding_dim))
+ layer = keras.layers.UnifiedLSTM(units, return_state=True, stateful=True)
+ outputs = layer(inputs)
+ state = outputs[1:]
+ assert len(state) == num_states
+ model = keras.models.Model(inputs, state[0])
+
+ inputs = np.random.random((num_samples, timesteps, embedding_dim))
+ state = model.predict(inputs)
+ self.assertAllClose(keras.backend.eval(layer.states[0]), state, atol=1e-4)
+
+ def test_state_reuse(self):
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+
+ inputs = keras.Input(batch_shape=(num_samples, timesteps, embedding_dim))
+ layer = keras.layers.UnifiedLSTM(
+ units, return_state=True, return_sequences=True)
+ outputs = layer(inputs)
+ output, state = outputs[0], outputs[1:]
+ output = keras.layers.UnifiedLSTM(units)(output, initial_state=state)
+ model = keras.models.Model(inputs, output)
+
+ inputs = np.random.random((num_samples, timesteps, embedding_dim))
+ model.predict(inputs)
+
+ def test_initial_states_as_other_inputs(self):
+ timesteps = 3
+ embedding_dim = 4
+ units = 3
+ num_samples = 2
+ num_states = 2
+ layer_class = keras.layers.UnifiedLSTM
+
+ # Test with Keras tensor
+ main_inputs = keras.Input((timesteps, embedding_dim))
+ initial_state = [keras.Input((units,)) for _ in range(num_states)]
+ inputs = [main_inputs] + initial_state
+
+ layer = layer_class(units)
+ output = layer(inputs)
+ assert initial_state[0] in layer._inbound_nodes[0].input_tensors
+
+ model = keras.models.Model(inputs, output)
+ model.compile(
+ loss='categorical_crossentropy',
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+
+ main_inputs = np.random.random((num_samples, timesteps, embedding_dim))
+ initial_state = [
+ np.random.random((num_samples, units)) for _ in range(num_states)
+ ]
+ targets = np.random.random((num_samples, units))
+ model.train_on_batch([main_inputs] + initial_state, targets)
+
+
+class LSTMLayerGraphOnlyTest(test.TestCase):
+
+ # Need session for test
@test_util.run_deprecated_v1
def test_unifiedLSTM(self):
input_shape = 10
@@ -71,7 +317,7 @@
num_classes=output_shape)
y_train = keras.utils.to_categorical(y_train, output_shape)
- layer = UnifiedLSTM(rnn_state_size, return_runtime=True)
+ layer = keras.layers.UnifiedLSTM(rnn_state_size, return_runtime=True)
inputs = array_ops.placeholder(
dtypes.float32, shape=(None, timestep, input_shape), name='inputs')
@@ -99,6 +345,7 @@
self.assertNotEqual(existing_loss, loss_value)
existing_loss = loss_value
+ # Need session for test
@test_util.run_deprecated_v1
def test_unifiedLSTM_with_cond(self):
# This test is to demonstrate the graph rewrite of grappler plugin under
@@ -119,7 +366,7 @@
num_classes=output_shape)
y_train = keras.utils.to_categorical(y_train, output_shape)
- layer = UnifiedLSTM(rnn_state_size, return_runtime=True)
+ layer = keras.layers.UnifiedLSTM(rnn_state_size, return_runtime=True)
inputs = array_ops.placeholder(
dtypes.float32, shape=(None, timestep, input_shape), name='inputs')
@@ -157,134 +404,32 @@
self.assertNotEqual(existing_loss, loss_value)
existing_loss = loss_value
- @parameterized.named_parameters(
- ('_non_tan_activation', 'relu', 0, False, True, None),
- ('_use_recurrent_dropout', 'tanh', 0.1, False, True, None),
- ('_unroll', 'tanh', 0, True, True, None),
- ('_not_use_bias', 'tanh', 0, False, False, None),
- ('_use_bias_regularizer', 'tanh', 0, False, True, 'l2')
- )
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_could_use_defun_backend(self, activation, recurrent_dropout,
- unroll, use_bias, bias_regularizer):
- layer = UnifiedLSTM(1,
- activation=activation,
- recurrent_dropout=recurrent_dropout,
- unroll=unroll,
- use_bias=use_bias,
- bias_regularizer=bias_regularizer)
- self.assertFalse(layer.could_use_cudnn)
-
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_unified_lstm_output_on_multiple_kernel(self):
- input_shape = 10
- rnn_state_size = 8
- timestep = 4
- batch = 100
-
- x_train = np.random.random((batch, timestep, input_shape))
-
- inputs = keras.layers.Input(
- shape=[timestep, input_shape], dtype=dtypes.float32)
- with test_util.device(use_gpu=False):
- # Note that CuDNN use 'sigmoid' as activation. Force the CPU
- # implementation to use 'sigmoid' so that it will generate same output as
- # CuDNN implementation.
- layer = UnifiedLSTM(rnn_state_size, recurrent_activation='sigmoid')
- output = layer(inputs)
- cpu_model = keras.models.Model(inputs, output)
- weights = cpu_model.get_weights()
- y_1 = cpu_model.predict(x_train)
-
- with test_util.device(use_gpu=True):
- layer = UnifiedLSTM(rnn_state_size, recurrent_activation='sigmoid')
- output = layer(inputs)
- gpu_model = keras.models.Model(inputs, output)
- gpu_model.set_weights(weights)
- y_2 = gpu_model.predict(x_train)
-
- with test_util.device(use_gpu=True):
- layer = keras.layers.LSTM(rnn_state_size, recurrent_activation='sigmoid')
- output = layer(inputs)
- canonical_model = keras.models.Model(inputs, output)
- # Remove the extra cudnn bias since canonical lstm will not use it.
- canonical_model.set_weights(weights[:3])
- y_3 = canonical_model.predict(x_train)
-
- self.assertAllClose(y_1, y_2)
- self.assertAllClose(y_2, y_3)
-
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_keras_model_with_lstm(self):
- input_shape = 10
- rnn_state_size = 8
- output_shape = 8
- timestep = 4
- batch = 100
- epoch = 10
-
- (x_train, y_train), _ = testing_utils.get_test_data(
- train_samples=batch,
- test_samples=0,
- input_shape=(timestep, input_shape),
- num_classes=output_shape)
- y_train = keras.utils.to_categorical(y_train, output_shape)
-
- layer = UnifiedLSTM(rnn_state_size)
-
- inputs = keras.layers.Input(
- shape=[timestep, input_shape], dtype=dtypes.float32)
-
- outputs = layer(inputs)
- model = keras.models.Model(inputs, outputs)
- model.compile('rmsprop', loss='mse')
- model.fit(x_train, y_train, epochs=epoch)
- model.evaluate(x_train, y_train)
-
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_return_sequences_LSTM(self):
- num_samples = 2
- timesteps = 3
+ # b/120919032
+ @test_util.run_deprecated_v1
+ def test_regularizers_LSTM(self):
embedding_dim = 4
- units = 2
- testing_utils.layer_test(
- UnifiedLSTM,
- kwargs={
- 'units': units,
- 'return_sequences': True
- },
- input_shape=(num_samples, timesteps, embedding_dim))
+ layer_class = keras.layers.UnifiedLSTM
+ layer = layer_class(
+ 5,
+ return_sequences=False,
+ weights=None,
+ input_shape=(None, embedding_dim),
+ kernel_regularizer=keras.regularizers.l1(0.01),
+ recurrent_regularizer=keras.regularizers.l1(0.01),
+ bias_regularizer='l2',
+ activity_regularizer='l1')
+ layer.build((None, None, 2))
+ self.assertEqual(len(layer.losses), 3)
+ x = keras.backend.variable(np.ones((2, 3, 2)))
+ layer(x)
+ self.assertEqual(len(layer.get_losses_for(x)), 1)
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_static_shape_inference_LSTM(self):
- # Github issue: 15165
- timesteps = 3
- embedding_dim = 4
- units = 2
- model = keras.models.Sequential()
- inputs = keras.layers.Dense(
- embedding_dim, input_shape=(timesteps, embedding_dim))
- model.add(inputs)
- layer = UnifiedLSTM(units, return_sequences=True)
- model.add(layer)
- outputs = model.layers[-1].output
- self.assertEqual(outputs.get_shape().as_list(), [None, timesteps, units])
+# TODO(scottzhu): Re-enable those tests in v2 mode once bugs attached are fixed.
+@test_util.run_v1_only
+class LSTMLayerV1OnlyTest(test.TestCase, parameterized.TestCase):
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_dynamic_behavior_LSTM(self):
- num_samples = 2
- timesteps = 3
- embedding_dim = 4
- units = 2
- layer = UnifiedLSTM(units, input_shape=(None, embedding_dim))
- model = keras.models.Sequential()
- model.add(layer)
- model.compile(gradient_descent.GradientDescentOptimizer(0.001), 'mse')
- x = np.random.random((num_samples, timesteps, embedding_dim))
- y = np.random.random((num_samples, units))
- model.train_on_batch(x, y)
-
+ # b/120911602
@test_util.run_in_graph_and_eager_modes(config=_config)
def test_dropout_LSTM(self):
num_samples = 2
@@ -292,7 +437,7 @@
embedding_dim = 4
units = 2
testing_utils.layer_test(
- UnifiedLSTM,
+ keras.layers.UnifiedLSTM,
kwargs={
'units': units,
'dropout': 0.1,
@@ -300,7 +445,48 @@
},
input_shape=(num_samples, timesteps, embedding_dim))
- @parameterized.parameters([0, 1, 2])
+ # b/120911602
+ def test_unified_lstm_feature_parity_with_canonical_lstm(self):
+ with context.eager_mode():
+ # Run this test under eager only due to b/120160788 for model.set_weights.
+ input_shape = 10
+ rnn_state_size = 8
+ timestep = 4
+ batch = 20
+
+ (x_train, y_train), _ = testing_utils.get_test_data(
+ train_samples=batch,
+ test_samples=0,
+ input_shape=(timestep, input_shape),
+ num_classes=rnn_state_size)
+ y_train = keras.utils.to_categorical(y_train, rnn_state_size)
+
+ inputs = keras.layers.Input(
+ shape=[timestep, input_shape], dtype=dtypes.float32)
+ lstm_layer = keras.layers.LSTM(rnn_state_size,
+ recurrent_activation='sigmoid')
+ output = lstm_layer(inputs)
+ lstm_model = keras.models.Model(inputs, output)
+ weights = lstm_model.get_weights()
+ y_1 = lstm_model.predict(x_train)
+ lstm_model.compile('rmsprop', 'mse')
+ lstm_model.fit(x_train, y_train)
+ y_2 = lstm_model.predict(x_train)
+
+ with test_util.device(use_gpu=True):
+ cudnn_layer = keras.layers.UnifiedLSTM(rnn_state_size)
+ cudnn_model = keras.models.Model(inputs, cudnn_layer(inputs))
+ cudnn_model.set_weights(weights)
+ y_3 = cudnn_model.predict(x_train)
+ cudnn_model.compile('rmsprop', 'mse')
+ cudnn_model.fit(x_train, y_train)
+ y_4 = cudnn_model.predict(x_train)
+
+ self.assertAllClose(y_1, y_3)
+ self.assertAllClose(y_2, y_4)
+
+ # b/120911602
+ @parameterized.named_parameters(('v0', 0), ('v1', 1), ('v2', 2))
@test_util.run_in_graph_and_eager_modes(config=_config)
def test_implementation_mode_LSTM(self, implementation_mode):
num_samples = 2
@@ -308,17 +494,14 @@
embedding_dim = 4
units = 2
testing_utils.layer_test(
- UnifiedLSTM,
+ keras.layers.UnifiedLSTM,
kwargs={
'units': units,
'implementation': implementation_mode
},
input_shape=(num_samples, timesteps, embedding_dim))
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_constraints_LSTM(self):
- embedding_dim = 4
- layer_class = UnifiedLSTM
+ layer_class = keras.layers.UnifiedLSTM
k_constraint = keras.constraints.max_norm(0.01)
r_constraint = keras.constraints.max_norm(0.01)
b_constraint = keras.constraints.max_norm(0.01)
@@ -335,9 +518,7 @@
self.assertEqual(layer.cell.recurrent_kernel.constraint, r_constraint)
self.assertEqual(layer.cell.bias.constraint, b_constraint)
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_with_masking_layer_LSTM(self):
- layer_class = UnifiedLSTM
+ layer_class = keras.layers.UnifiedLSTM
inputs = np.random.random((2, 3, 4))
targets = np.abs(np.random.random((2, 3, 5)))
targets /= targets.sum(axis=-1, keepdims=True)
@@ -349,19 +530,7 @@
optimizer=gradient_descent.GradientDescentOptimizer(0.01))
model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_stacking_LSTM(self):
- inputs = np.random.random((2, 3, 4))
- targets = np.abs(np.random.random((2, 3, 5)))
- targets /= targets.sum(axis=-1, keepdims=True)
- model = keras.models.Sequential()
- model.add(UnifiedLSTM(10, return_sequences=True, unroll=False))
- model.add(UnifiedLSTM(5, return_sequences=True, unroll=False))
- model.compile(
- loss='categorical_crossentropy',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
- model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
-
+ # b/120911602
@test_util.run_in_graph_and_eager_modes(config=_config)
def test_masking_with_stacking_LSTM(self):
inputs = np.random.random((2, 3, 4))
@@ -369,286 +538,245 @@
targets /= targets.sum(axis=-1, keepdims=True)
model = keras.models.Sequential()
model.add(keras.layers.Masking(input_shape=(3, 4)))
- model.add(UnifiedLSTM(10, return_sequences=True, unroll=False))
- model.add(UnifiedLSTM(5, return_sequences=True, unroll=False))
+ model.add(keras.layers.UnifiedLSTM(10, return_sequences=True, unroll=False))
+ model.add(keras.layers.UnifiedLSTM(5, return_sequences=True, unroll=False))
model.compile(
loss='categorical_crossentropy',
optimizer=gradient_descent.GradientDescentOptimizer(0.01))
model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
+ # b/120911602
+ @parameterized.named_parameters(
+ # test_name, time_major, go_backwards
+ ('normal', False, False),
+ ('time_major', True, False),
+ ('go_backwards', False, True),
+ ('both', True, True),
+ )
@test_util.run_in_graph_and_eager_modes(config=_config)
- def test_from_config_LSTM(self):
- layer_class = UnifiedLSTM
- for stateful in (False, True):
- l1 = layer_class(units=1, stateful=stateful)
- l2 = layer_class.from_config(l1.get_config())
- assert l1.get_config() == l2.get_config()
+ def test_time_major_and_go_backward(self, time_major, go_backwards):
+ input_shape = 10
+ rnn_state_size = 8
+ timestep = 4
+ batch = 100
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_specify_initial_state_keras_tensor(self):
- num_states = 2
- timesteps = 3
- embedding_dim = 4
- units = 3
- num_samples = 2
+ x_train = np.random.random((batch, timestep, input_shape))
- # Test with Keras tensor
- inputs = keras.Input((timesteps, embedding_dim))
- initial_state = [keras.Input((units,)) for _ in range(num_states)]
- layer = UnifiedLSTM(units)
- if len(initial_state) == 1:
- output = layer(inputs, initial_state=initial_state[0])
- else:
- output = layer(inputs, initial_state=initial_state)
- assert initial_state[0] in layer._inbound_nodes[0].input_tensors
+ def build_model(layer_cls):
+ inputs = keras.layers.Input(
+ shape=[timestep, input_shape], dtype=dtypes.float32)
+ layer = layer_cls(rnn_state_size,
+ recurrent_activation='sigmoid',
+ time_major=time_major,
+ return_sequences=True,
+ go_backwards=go_backwards)
+ if time_major:
+ converted_input = keras.layers.Lambda(
+ lambda t: array_ops.transpose(t, [1, 0, 2]))(inputs)
+ outputs = layer(converted_input)
+ outputs = keras.layers.Lambda(
+ lambda t: array_ops.transpose(t, [1, 0, 2]))(outputs)
+ else:
+ outputs = layer(inputs)
+ return keras.models.Model(inputs, outputs)
- model = keras.models.Model([inputs] + initial_state, output)
- model.compile(
- loss='categorical_crossentropy',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+ lstm_model = build_model(keras.layers.LSTM)
+ y_ref = lstm_model.predict(x_train)
+ weights = lstm_model.get_weights()
- inputs = np.random.random((num_samples, timesteps, embedding_dim))
- initial_state = [
- np.random.random((num_samples, units)) for _ in range(num_states)
- ]
- targets = np.random.random((num_samples, units))
- model.train_on_batch([inputs] + initial_state, targets)
+ unified_lstm_model = build_model(keras.layers.UnifiedLSTM)
+ unified_lstm_model.set_weights(weights)
+ y = unified_lstm_model.predict(x_train)
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def DISABLED_test_specify_initial_state_non_keras_tensor(self):
- num_states = 2
- timesteps = 3
- embedding_dim = 4
- units = 3
- num_samples = 2
+ self.assertAllClose(y, y_ref)
- # Test with non-Keras tensor
- inputs = keras.Input((timesteps, embedding_dim))
- initial_state = [
- keras.backend.random_normal_variable((num_samples, units), 0, 1)
- for _ in range(num_states)
- ]
- layer = UnifiedLSTM(units)
- output = layer(inputs, initial_state=initial_state)
+ input_shape = 10
+ rnn_state_size = 8
+ output_shape = 8
+ timestep = 4
+ batch = 100
+ epoch = 10
- model = keras.models.Model(inputs, output)
- model.compile(
- loss='categorical_crossentropy',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
+ (x_train, y_train), _ = testing_utils.get_test_data(
+ train_samples=batch,
+ test_samples=0,
+ input_shape=(timestep, input_shape),
+ num_classes=output_shape)
+ y_train = keras.utils.to_categorical(y_train, output_shape)
- inputs = np.random.random((num_samples, timesteps, embedding_dim))
- targets = np.random.random((num_samples, units))
- model.train_on_batch(inputs, targets)
+ layer = keras.layers.UnifiedLSTM(rnn_state_size)
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_reset_states_with_values(self):
- num_states = 2
- timesteps = 3
- embedding_dim = 4
- units = 3
- num_samples = 2
+ inputs = keras.layers.Input(
+ shape=[timestep, input_shape], dtype=dtypes.float32)
- layer = UnifiedLSTM(units, stateful=True)
- layer.build((num_samples, timesteps, embedding_dim))
- layer.reset_states()
- assert len(layer.states) == num_states
- assert layer.states[0] is not None
- self.assertAllClose(
- keras.backend.eval(layer.states[0]),
- np.zeros(keras.backend.int_shape(layer.states[0])),
- atol=1e-4)
- state_shapes = [keras.backend.int_shape(state) for state in layer.states]
- values = [np.ones(shape) for shape in state_shapes]
- if len(values) == 1:
- values = values[0]
- layer.reset_states(values)
- self.assertAllClose(
- keras.backend.eval(layer.states[0]),
- np.ones(keras.backend.int_shape(layer.states[0])),
- atol=1e-4)
-
- # Test with invalid data
- with self.assertRaises(ValueError):
- layer.reset_states([1] * (len(layer.states) + 1))
-
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_specify_state_with_masking(self):
- num_states = 2
- timesteps = 3
- embedding_dim = 4
- units = 3
- num_samples = 2
-
- inputs = keras.Input((timesteps, embedding_dim))
- _ = keras.layers.Masking()(inputs)
- initial_state = [keras.Input((units,)) for _ in range(num_states)]
- output = UnifiedLSTM(units)(inputs, initial_state=initial_state)
-
- model = keras.models.Model([inputs] + initial_state, output)
- model.compile(
- loss='categorical_crossentropy',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
-
- inputs = np.random.random((num_samples, timesteps, embedding_dim))
- initial_state = [
- np.random.random((num_samples, units)) for _ in range(num_states)
- ]
- targets = np.random.random((num_samples, units))
- model.train_on_batch([inputs] + initial_state, targets)
-
- @test_util.run_in_graph_and_eager_modes(config=_config)
- def test_return_state(self):
- num_states = 2
- timesteps = 3
- embedding_dim = 4
- units = 3
- num_samples = 2
-
- inputs = keras.Input(batch_shape=(num_samples, timesteps, embedding_dim))
- layer = UnifiedLSTM(units, return_state=True, stateful=True)
outputs = layer(inputs)
- state = outputs[1:]
- assert len(state) == num_states
- model = keras.models.Model(inputs, state[0])
+ model = keras.models.Model(inputs, outputs)
+ model.compile('rmsprop', loss='mse')
+ model.fit(x_train, y_train, epochs=epoch)
+ model.evaluate(x_train, y_train)
+ model.predict(x_train)
- inputs = np.random.random((num_samples, timesteps, embedding_dim))
- state = model.predict(inputs)
- self.assertAllClose(keras.backend.eval(layer.states[0]), state, atol=1e-4)
+ # b/120911602
+ @parameterized.named_parameters(
+ # test_name, use_bias, bias_initializer, activation
+ ('normal', True, 'zeros'),
+ ('no_bias', False, 'zeros'),
+ ('random_bias', True, 'random_uniform'),
+ )
+ @test_util.run_in_graph_and_eager_modes(config=_config)
+ def test_unified_lstm_model_save_load(self, use_bias, bias_initializer):
+ temp_dir = self.get_temp_dir()
+ self.addCleanup(shutil.rmtree, temp_dir)
+ h5_path = os.path.join(temp_dir, 'test.h5')
+
+ batch = 10
+ timestep = 3
+ input_dim = 5
+ units = 2
+
+ x = np.random.random((batch, timestep, input_dim))
+
+ def build_model():
+ inputs = keras.layers.Input(
+ shape=[timestep, input_dim], dtype=dtypes.float32)
+ layer = keras.layers.UnifiedLSTM(
+ units,
+ use_bias=use_bias,
+ bias_initializer=bias_initializer)
+ output = layer(inputs)
+ return keras.models.Model(inputs, output), layer
+
+ model, layer = build_model()
+ y_ref = model.predict(x)
+ model.save_weights(h5_path)
+
+ cloned_model, new_layer = build_model()
+ cloned_model.load_weights(h5_path)
+ y = cloned_model.predict(x)
+
+ self.assertAllClose(y, y_ref)
+ self.assertAllClose(layer.get_weights(), new_layer.get_weights())
+
+ # b/120911602
+ @test_util.run_in_graph_and_eager_modes(config=_config)
+ def test_unified_lstm_output_on_multiple_kernel(self):
+ input_shape = 10
+ rnn_state_size = 8
+ timestep = 4
+ batch = 100
+
+ x_train = np.random.random((batch, timestep, input_shape))
+
+ inputs = keras.layers.Input(
+ shape=[timestep, input_shape], dtype=dtypes.float32)
+ with test_util.device(use_gpu=False):
+ layer = keras.layers.UnifiedLSTM(rnn_state_size)
+ output = layer(inputs)
+ cpu_model = keras.models.Model(inputs, output)
+ weights = cpu_model.get_weights()
+ y_1 = cpu_model.predict(x_train)
+
+ with test_util.device(use_gpu=True):
+ layer = keras.layers.UnifiedLSTM(rnn_state_size)
+ output = layer(inputs)
+ gpu_model = keras.models.Model(inputs, output)
+ gpu_model.set_weights(weights)
+ y_2 = gpu_model.predict(x_train)
+
+ # Note that CuDNN uses 'sigmoid' as activation, so the unified LSTM uses
+ # 'sigmoid' as default. Construct the canonical LSTM with sigmoid to achieve
+ # the same output.
+ with test_util.device(use_gpu=True):
+ layer = keras.layers.LSTM(rnn_state_size, recurrent_activation='sigmoid')
+ output = layer(inputs)
+ canonical_model = keras.models.Model(inputs, output)
+ # Remove the extra cudnn bias since canonical lstm will not use it.
+ canonical_model.set_weights(weights[:3])
+ y_3 = canonical_model.predict(x_train)
+
+ self.assertAllClose(y_1, y_2)
+ self.assertAllClose(y_2, y_3)
@test_util.run_in_graph_and_eager_modes(config=_config)
- def test_state_reuse(self):
+ def test_return_sequences_LSTM(self):
+ num_samples = 2
timesteps = 3
embedding_dim = 4
- units = 3
- num_samples = 2
+ units = 2
+ testing_utils.layer_test(
+ keras.layers.UnifiedLSTM,
+ kwargs={
+ 'units': units,
+ 'return_sequences': True
+ },
+ input_shape=(num_samples, timesteps, embedding_dim))
- inputs = keras.Input(batch_shape=(num_samples, timesteps, embedding_dim))
- layer = UnifiedLSTM(units, return_state=True, return_sequences=True)
- outputs = layer(inputs)
- output, state = outputs[0], outputs[1:]
- output = UnifiedLSTM(units)(output, initial_state=state)
- model = keras.models.Model(inputs, output)
-
- inputs = np.random.random((num_samples, timesteps, embedding_dim))
- model.predict(inputs)
-
+ # b/120911602
@test_util.run_in_graph_and_eager_modes(config=_config)
- def test_initial_states_as_other_inputs(self):
- timesteps = 3
- embedding_dim = 4
- units = 3
- num_samples = 2
- num_states = 2
- layer_class = UnifiedLSTM
-
- # Test with Keras tensor
- main_inputs = keras.Input((timesteps, embedding_dim))
- initial_state = [keras.Input((units,)) for _ in range(num_states)]
- inputs = [main_inputs] + initial_state
-
- layer = layer_class(units)
- output = layer(inputs)
- assert initial_state[0] in layer._inbound_nodes[0].input_tensors
-
- model = keras.models.Model(inputs, output)
- model.compile(
- loss='categorical_crossentropy',
- optimizer=gradient_descent.GradientDescentOptimizer(0.01))
-
- main_inputs = np.random.random((num_samples, timesteps, embedding_dim))
- initial_state = [
- np.random.random((num_samples, units)) for _ in range(num_states)
- ]
- targets = np.random.random((num_samples, units))
- model.train_on_batch([main_inputs] + initial_state, targets)
-
-
-class LSTMLayerGraphOnlyTest(test.TestCase):
-
def test_statefulness_LSTM(self):
num_samples = 2
timesteps = 3
embedding_dim = 4
units = 2
- layer_class = UnifiedLSTM
- with self.cached_session(config=_config):
- model = keras.models.Sequential()
- model.add(
- keras.layers.Embedding(
- 4,
- embedding_dim,
- mask_zero=True,
- input_length=timesteps,
- batch_input_shape=(num_samples, timesteps)))
- layer = layer_class(
- units, return_sequences=False, stateful=True, weights=None)
- model.add(layer)
- model.compile(
- optimizer=gradient_descent.GradientDescentOptimizer(0.01), loss='mse')
- out1 = model.predict(np.ones((num_samples, timesteps)))
- self.assertEqual(out1.shape, (num_samples, units))
+ layer_class = keras.layers.UnifiedLSTM
+ model = keras.models.Sequential()
+ model.add(
+ keras.layers.Embedding(
+ 4,
+ embedding_dim,
+ mask_zero=True,
+ input_length=timesteps,
+ batch_input_shape=(num_samples, timesteps)))
+ layer = layer_class(
+ units, return_sequences=False, stateful=True, weights=None)
+ model.add(layer)
+ model.compile(
+ optimizer=gradient_descent.GradientDescentOptimizer(0.01), loss='mse')
+ out1 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertEqual(out1.shape, (num_samples, units))
- # train once so that the states change
- model.train_on_batch(
- np.ones((num_samples, timesteps)), np.ones((num_samples, units)))
- out2 = model.predict(np.ones((num_samples, timesteps)))
+ # train once so that the states change
+ model.train_on_batch(
+ np.ones((num_samples, timesteps)), np.ones((num_samples, units)))
+ out2 = model.predict(np.ones((num_samples, timesteps)))
- # if the state is not reset, output should be different
- self.assertNotEqual(out1.max(), out2.max())
+ # if the state is not reset, output should be different
+ self.assertNotEqual(out1.max(), out2.max())
- # check that output changes after states are reset
- # (even though the model itself didn't change)
- layer.reset_states()
- out3 = model.predict(np.ones((num_samples, timesteps)))
- self.assertNotEqual(out2.max(), out3.max())
+ # check that output changes after states are reset
+ # (even though the model itself didn't change)
+ layer.reset_states()
+ out3 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertNotEqual(out2.max(), out3.max())
- # check that container-level reset_states() works
- model.reset_states()
- out4 = model.predict(np.ones((num_samples, timesteps)))
- self.assertAllClose(out3, out4, atol=1e-5)
+ # check that container-level reset_states() works
+ model.reset_states()
+ out4 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertAllClose(out3, out4, atol=1e-5)
- # check that the call to `predict` updated the states
- out5 = model.predict(np.ones((num_samples, timesteps)))
- self.assertNotEqual(out4.max(), out5.max())
+ # check that the call to `predict` updated the states
+ out5 = model.predict(np.ones((num_samples, timesteps)))
+ self.assertNotEqual(out4.max(), out5.max())
- # Check masking
- layer.reset_states()
+ # Check masking
+ layer.reset_states()
- left_padded_input = np.ones((num_samples, timesteps))
- left_padded_input[0, :1] = 0
- left_padded_input[1, :2] = 0
- out6 = model.predict(left_padded_input)
+ left_padded_input = np.ones((num_samples, timesteps))
+ left_padded_input[0, :1] = 0
+ left_padded_input[1, :2] = 0
+ out6 = model.predict(left_padded_input)
- layer.reset_states()
+ layer.reset_states()
- right_padded_input = np.ones((num_samples, timesteps))
- right_padded_input[0, -1:] = 0
- right_padded_input[1, -2:] = 0
- out7 = model.predict(right_padded_input)
+ right_padded_input = np.ones((num_samples, timesteps))
+ right_padded_input[0, -1:] = 0
+ right_padded_input[1, -2:] = 0
+ out7 = model.predict(right_padded_input)
- self.assertAllClose(out7, out6, atol=1e-5)
-
- def test_regularizers_LSTM(self):
- embedding_dim = 4
- layer_class = UnifiedLSTM
- with self.cached_session(config=_config):
- layer = layer_class(
- 5,
- return_sequences=False,
- weights=None,
- input_shape=(None, embedding_dim),
- kernel_regularizer=keras.regularizers.l1(0.01),
- recurrent_regularizer=keras.regularizers.l1(0.01),
- bias_regularizer='l2',
- activity_regularizer='l1')
- layer.build((None, None, 2))
- self.assertEqual(len(layer.losses), 3)
- x = keras.backend.variable(np.ones((2, 3, 2)))
- layer(x)
- self.assertEqual(len(layer.get_losses_for(x)), 1)
+ self.assertAllClose(out7, out6, atol=1e-5)
-class UnifiedLSTMPerformanceTest(test.TestCase):
+class UnifiedLSTMPerformanceTest(test.Benchmark):
def _measure_performance(self, test_config, model, x_train, y_train):
batch = test_config['batch']
@@ -668,7 +796,7 @@
rnn_state_size = test_config['rnn_state_size']
timestep = test_config['timestep']
- cudnn_lstm_layer = CuDNNLSTM(rnn_state_size)
+ cudnn_lstm_layer = keras.layers.CuDNNLSTM(rnn_state_size)
inputs = keras.layers.Input(
shape=[timestep, input_shape], dtype=dtypes.float32)
@@ -689,7 +817,7 @@
rnn_state_size = test_config['rnn_state_size']
timestep = test_config['timestep']
- layer = UnifiedLSTM(rnn_state_size)
+ layer = keras.layers.UnifiedLSTM(rnn_state_size)
inputs = keras.layers.Input(
shape=[timestep, input_shape], dtype=dtypes.float32)
@@ -724,11 +852,11 @@
'Normal LSTM', sec_per_epoch)
return sec_per_epoch
- @test_util.run_in_graph_and_eager_modes(config=_config, use_gpu=True)
- def test_performance_with_standard_cudnn_impl(self):
+ def _benchmark_performance_with_standard_cudnn_impl(self):
if not test.is_gpu_available():
self.skipTest('performance test will only run on GPU')
+ mode = 'eager' if context.executing_eagerly() else 'graph'
batch = 64
num_batch = 10
test_config = {
@@ -748,34 +876,42 @@
num_classes=test_config['output_shape'])
y_train = keras.utils.to_categorical(y_train, test_config['output_shape'])
- cudnn_duration = self._time_performance_run_cudnn_lstm(
+ cudnn_sec_per_epoch = self._time_performance_run_cudnn_lstm(
test_config, x_train, y_train)
- unified_lstm_gpu_duration = self._time_performance_run_unifed_lstm_gpu(
+ unified_lstm_sec_per_epoch = self._time_performance_run_unifed_lstm_gpu(
test_config, x_train, y_train)
- normal_lstm_duration = self._time_performance_run_normal_lstm(
+ normal_lstm_sec_per_epoch = self._time_performance_run_normal_lstm(
test_config, x_train, y_train)
- cudnn_vs_unified = cudnn_duration / unified_lstm_gpu_duration
- unified_vs_normal = normal_lstm_duration / unified_lstm_gpu_duration
+ cudnn_vs_unified = cudnn_sec_per_epoch / unified_lstm_sec_per_epoch
+ unified_vs_normal = normal_lstm_sec_per_epoch / unified_lstm_sec_per_epoch
- # TODO(scottzhu): reeanble the test after moving it to benchmark test suite.
- # The current test has performance flakiness issue.
+ self.report_benchmark(name='keras_cudnn_lstm_' + mode,
+ wall_time=cudnn_sec_per_epoch,
+ iters=test_config['epoch'],
+ extras=test_config)
+ self.report_benchmark(name='keras_unified_lstm_' + mode,
+ wall_time=unified_lstm_sec_per_epoch,
+ iters=test_config['epoch'],
+ extras=test_config)
+ self.report_benchmark(name='keras_canonical_lstm_' + mode,
+ wall_time=normal_lstm_sec_per_epoch,
+ iters=test_config['epoch'],
+ extras=test_config)
+
logging.info('Expect the performance of Unified LSTM is within 80% of '
'CuDNN LSTM, got {0:.2f}%'.format(cudnn_vs_unified * 100))
logging.info('Expect the performance of Unified LSTM is more than 5 times'
' of normal LSTM, got {0:.2f}'.format(unified_vs_normal))
- # Assert the performance diff should be within 80% of the native cudnn.
- # self.assertGreaterEqual(
- # cudnn_vs_unified, 0.80,
- # 'Expect the performance of Unified LSTM is within 80% of CuDNN LSTM, '
- # 'but got {0:.2f}%'.format(cudnn_vs_unified * 100))
- # # Assert the performance diff between CPU impl and GPU impl should be more
- # # than 5 times.
- # self.assertGreaterEqual(
- # unified_vs_normal, 5,
- # 'Expect the performance of Unified LSTM is more than 5 times of '
- # 'normal LSTM, but got {0:.2f}'.format(unified_vs_normal))
+ def benchmark_performance_graph(self):
+ with context.graph_mode(), session_lib.Session(config=_config):
+ self._benchmark_performance_with_standard_cudnn_impl()
+
+ def benchmark_performance_eager(self):
+ with context.eager_mode():
+ self._benchmark_performance_with_standard_cudnn_impl()
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/layers/wrappers_test.py b/tensorflow/python/keras/layers/wrappers_test.py
index b919641..727f33d 100644
--- a/tensorflow/python/keras/layers/wrappers_test.py
+++ b/tensorflow/python/keras/layers/wrappers_test.py
@@ -165,6 +165,7 @@
y = model.predict(np.random.random((10, 3, 2)))
self.assertAllClose(np.mean(y), 0., atol=1e-1, rtol=1e-1)
+ @tf_test_util.run_v1_only('b/120545219')
def test_TimeDistributed_batchnorm(self):
with self.cached_session():
# test that wrapped BN updates still work.
@@ -187,6 +188,7 @@
# Verify input_map has one mapping from inputs to reshaped inputs.
self.assertEqual(len(td._input_map.keys()), 1)
+ @tf_test_util.run_v1_only('b/120545219')
def test_TimeDistributed_trainable(self):
# test layers that need learning_phase to be set
x = keras.layers.Input(shape=(3, 2))
@@ -201,7 +203,7 @@
assert len(layer.updates) == 2
assert len(layer.trainable_weights) == 2
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_TimeDistributed_with_masked_embedding_and_unspecified_shape(self):
with self.cached_session():
# test with unspecified shape and Embeddings with mask_zero
@@ -234,7 +236,7 @@
self.assertAllEqual(mask_outputs_val[i], ref_mask_val[i])
self.assertIs(mask_outputs[-1], None) # final layer
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_TimeDistributed_with_masking_layer(self):
with self.cached_session():
# test with Masking layer
@@ -377,7 +379,7 @@
model.compile(loss='mse', optimizer='sgd')
model.fit(x, y, epochs=1, batch_size=1)
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_Bidirectional_merged_value(self):
rnn = keras.layers.LSTM
samples = 2
@@ -508,7 +510,7 @@
layer.trainable = True
assert len(layer.trainable_weights) == 6
- @tf_test_util.run_deprecated_v1
+ @tf_test_util.run_v1_only('b/120545219')
def test_Bidirectional_updates(self):
with self.cached_session():
x = keras.layers.Input(shape=(3, 2))
diff --git a/tensorflow/python/keras/metrics.py b/tensorflow/python/keras/metrics.py
index 3c2682e..c8ccb7f 100644
--- a/tensorflow/python/keras/metrics.py
+++ b/tensorflow/python/keras/metrics.py
@@ -28,6 +28,7 @@
import numpy as np
import six
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.eager import context
from tensorflow.python.eager import function
from tensorflow.python.framework import dtypes
@@ -60,7 +61,6 @@
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables as tf_variables
from tensorflow.python.ops import weights_broadcast_ops
-from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.util import tf_decorator
from tensorflow.python.util.tf_export import tf_export
from tensorflow.tools.docs import doc_controls
@@ -171,12 +171,18 @@
def _assert_thresholds_range(thresholds):
- invalid_thresholds = [t for t in thresholds if t < 0 or t > 1]
- if any(invalid_thresholds):
+ invalid_thresholds = [t for t in thresholds if t is None or t < 0 or t > 1]
+ if invalid_thresholds:
raise ValueError('Threshold values must be in [0, 1]. Invalid values: {}'
.format(invalid_thresholds))
+def _parse_init_thresholds(thresholds, default_threshold=0.5):
+ thresholds = to_list(default_threshold if thresholds is None else thresholds)
+ _assert_thresholds_range(thresholds)
+ return thresholds
+
+
def _update_confusion_matrix_variables(variables_to_update,
y_true,
y_pred,
@@ -511,7 +517,7 @@
### End: For use by subclasses ###
-@tf_export('metrics.Mean', 'keras.metrics.Mean')
+@tf_export('keras.metrics.Mean')
class Mean(Metric):
"""Computes the (weighted) mean of the given values.
@@ -528,7 +534,7 @@
Usage:
```python
- m = tf.metrics.Mean()
+ m = tf.keras.metrics.Mean()
m.update_state([1, 3, 5, 7])
print('Final result: ', m.result().numpy()) # Final result: 4.0
```
@@ -537,7 +543,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.add_metric(metrics_module.Mean(name='mean_1')(outputs))
+ model.add_metric(tf.keras.metrics.Mean(name='mean_1')(outputs))
model.compile('sgd', loss='mse')
```
"""
@@ -651,7 +657,7 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('metrics.Accuracy', 'keras.metrics.Accuracy')
+@tf_export('keras.metrics.Accuracy')
class Accuracy(MeanMetricWrapper):
"""Calculates how often predictions matches labels.
@@ -670,7 +676,7 @@
Usage:
```python
- m = tf.metrics.Accuracy()
+ m = tf.keras.metrics.Accuracy()
m.update_state([1, 2, 3, 4], [0, 2, 3, 4])
print('Final result: ', m.result().numpy()) # Final result: 0.75
```
@@ -679,7 +685,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.Accuracy()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.Accuracy()])
```
"""
@@ -693,7 +699,7 @@
return super(Accuracy, cls).from_config(config)
-@tf_export('metrics.BinaryAccuracy', 'keras.metrics.BinaryAccuracy')
+@tf_export('keras.metrics.BinaryAccuracy')
class BinaryAccuracy(MeanMetricWrapper):
"""Calculates how often predictions matches labels.
@@ -712,7 +718,7 @@
Usage:
```python
- m = tf.metrics.BinaryAccuracy()
+ m = tf.keras.metrics.BinaryAccuracy()
m.update_state([1, 1, 0, 0], [0.98, 1, 0, 0.6])
print('Final result: ', m.result().numpy()) # Final result: 0.75
```
@@ -721,7 +727,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.BinaryAccuracy()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.BinaryAccuracy()])
```
"""
@@ -744,8 +750,7 @@
return super(BinaryAccuracy, cls).from_config(config)
-@tf_export(
- 'metrics.CategoricalAccuracy', 'keras.metrics.CategoricalAccuracy')
+@tf_export('keras.metrics.CategoricalAccuracy')
class CategoricalAccuracy(MeanMetricWrapper):
"""Calculates how often predictions matches labels.
@@ -768,7 +773,7 @@
Usage:
```python
- m = tf.metrics.CategoricalAccuracy()
+ m = tf.keras.metrics.CategoricalAccuracy()
m.update_state([[0, 0, 1], [0, 1, 0]], [[0.1, 0.9, 0.8], [0.05, 0.95, 0]])
print('Final result: ', m.result().numpy()) # Final result: 0.5
```
@@ -777,7 +782,10 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.CategoricalAccuracy()])
+ model.compile(
+ 'sgd',
+ loss='mse',
+ metrics=[tf.keras.metrics.CategoricalAccuracy()])
```
"""
@@ -798,9 +806,7 @@
return super(CategoricalAccuracy, cls).from_config(config)
-@tf_export(
- 'metrics.SparseCategoricalAccuracy',
- 'keras.metrics.SparseCategoricalAccuracy')
+@tf_export('keras.metrics.SparseCategoricalAccuracy')
class SparseCategoricalAccuracy(MeanMetricWrapper):
"""Calculates how often predictions matches integer labels.
@@ -820,7 +826,7 @@
Usage:
```python
- m = tf.metrics.SparseCategoricalAccuracy()
+ m = tf.keras.metrics.SparseCategoricalAccuracy()
m.update_state([[2], [1]], [[0.1, 0.9, 0.8], [0.05, 0.95, 0]])
print('Final result: ', m.result().numpy()) # Final result: 0.5
```
@@ -832,7 +838,7 @@
model.compile(
'sgd',
loss='mse',
- metrics=[tf.metrics.SparseCategoricalAccuracy()])
+ metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
```
"""
@@ -869,12 +875,11 @@
"""
super(_ConfusionMatrixConditionCount, self).__init__(name=name, dtype=dtype)
self._confusion_matrix_cond = confusion_matrix_cond
- self.thresholds = 0.5 if thresholds is None else thresholds
- thresholds = to_list(thresholds)
- _assert_thresholds_range(thresholds)
+ self.thresholds = _parse_init_thresholds(
+ thresholds, default_threshold=0.5)
self.accumulator = self.add_weight(
'accumulator',
- shape=(len(thresholds),),
+ shape=(len(self.thresholds),),
initializer=init_ops.zeros_initializer)
def update_state(self, y_true, y_pred, sample_weight=None):
@@ -895,10 +900,10 @@
}, y_true, y_pred, self.thresholds, sample_weight)
def result(self):
- if isinstance(self.thresholds, (list, tuple)):
- result = self.accumulator
- else:
+ if len(self.thresholds) == 1:
result = self.accumulator[0]
+ else:
+ result = self.accumulator
return ops.convert_to_tensor(result)
def reset_states(self):
@@ -907,7 +912,7 @@
K.set_value(v, np.zeros((num_thresholds,)))
-@tf_export('metrics.FalsePositives', 'keras.metrics.FalsePositives')
+@tf_export('keras.metrics.FalsePositives')
class FalsePositives(_ConfusionMatrixConditionCount):
"""Calculates the number of false positives.
@@ -925,7 +930,7 @@
Usage:
```python
- m = tf.metrics.FalsePositives()
+ m = tf.keras.metrics.FalsePositives()
m.update_state([0, 1, 0, 0], [0, 0, 1, 1])
print('Final result: ', m.result().numpy()) # Final result: 2
```
@@ -934,7 +939,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.FalsePositives()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.FalsePositives()])
```
"""
@@ -957,7 +962,7 @@
dtype=dtype)
-@tf_export('metrics.FalseNegatives', 'keras.metrics.FalseNegatives')
+@tf_export('keras.metrics.FalseNegatives')
class FalseNegatives(_ConfusionMatrixConditionCount):
"""Calculates the number of false negatives.
@@ -975,7 +980,7 @@
Usage:
```python
- m = tf.metrics.FalseNegatives()
+ m = tf.keras.metrics.FalseNegatives()
m.update_state([0, 1, 1, 1], [0, 1, 0, 0])
print('Final result: ', m.result().numpy()) # Final result: 2
```
@@ -984,7 +989,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.FalseNegatives()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.FalseNegatives()])
```
"""
@@ -1007,7 +1012,7 @@
dtype=dtype)
-@tf_export('metrics.TrueNegatives', 'keras.metrics.TrueNegatives')
+@tf_export('keras.metrics.TrueNegatives')
class TrueNegatives(_ConfusionMatrixConditionCount):
"""Calculates the number of true negatives.
@@ -1025,7 +1030,7 @@
Usage:
```python
- m = tf.metrics.TrueNegatives()
+ m = tf.keras.metrics.TrueNegatives()
m.update_state([0, 1, 0, 0], [1, 1, 0, 0])
print('Final result: ', m.result().numpy()) # Final result: 2
```
@@ -1034,7 +1039,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.TrueNegatives()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.TrueNegatives()])
```
"""
@@ -1057,7 +1062,7 @@
dtype=dtype)
-@tf_export('metrics.TruePositives', 'keras.metrics.TruePositives')
+@tf_export('keras.metrics.TruePositives')
class TruePositives(_ConfusionMatrixConditionCount):
"""Calculates the number of true positives.
@@ -1075,7 +1080,7 @@
Usage:
```python
- m = tf.metrics.TruePositives()
+ m = tf.keras.metrics.TruePositives()
m.update_state([0, 1, 1, 1], [1, 0, 1, 1])
print('Final result: ', m.result().numpy()) # Final result: 2
```
@@ -1084,7 +1089,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.TruePositives()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.TruePositives()])
```
"""
@@ -1107,7 +1112,7 @@
dtype=dtype)
-@tf_export('metrics.Precision', 'keras.metrics.Precision')
+@tf_export('keras.metrics.Precision')
class Precision(Metric):
"""Computes the precision of the predictions with respect to the labels.
@@ -1126,7 +1131,7 @@
Usage:
```python
- m = tf.metrics.Precision()
+ m = tf.keras.metrics.Precision()
m.update_state([0, 1, 1, 1], [1, 0, 1, 1])
print('Final result: ', m.result().numpy()) # Final result: 0.66
```
@@ -1135,7 +1140,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.Precision()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.Precision()])
```
"""
@@ -1152,16 +1157,15 @@
dtype: (Optional) data type of the metric result.
"""
super(Precision, self).__init__(name=name, dtype=dtype)
- self.thresholds = 0.5 if thresholds is None else thresholds
- thresholds = to_list(thresholds)
- _assert_thresholds_range(thresholds)
+ self.thresholds = _parse_init_thresholds(
+ thresholds, default_threshold=0.5)
self.tp = self.add_weight(
'true_positives',
- shape=(len(thresholds),),
+ shape=(len(self.thresholds),),
initializer=init_ops.zeros_initializer)
self.fp = self.add_weight(
'false_positives',
- shape=(len(thresholds),),
+ shape=(len(self.thresholds),),
initializer=init_ops.zeros_initializer)
def update_state(self, y_true, y_pred, sample_weight=None):
@@ -1184,7 +1188,7 @@
def result(self):
result = math_ops.div_no_nan(self.tp, self.tp + self.fp)
- return result if isinstance(self.thresholds, (list, tuple)) else result[0]
+ return result[0] if len(self.thresholds) == 1 else result
def reset_states(self):
num_thresholds = len(to_list(self.thresholds))
@@ -1192,7 +1196,7 @@
K.set_value(v, np.zeros((num_thresholds,)))
-@tf_export('metrics.Recall', 'keras.metrics.Recall')
+@tf_export('keras.metrics.Recall')
class Recall(Metric):
"""Computes the recall of the predictions with respect to the labels.
@@ -1211,7 +1215,7 @@
Usage:
```python
- m = tf.metrics.Recall()
+ m = tf.keras.metrics.Recall()
m.update_state([0, 1, 1, 1], [1, 0, 1, 1])
print('Final result: ', m.result().numpy()) # Final result: 0.66
```
@@ -1220,7 +1224,7 @@
```python
model = keras.models.Model(inputs, outputs)
- model.compile('sgd', loss='mse', metrics=[tf.metrics.Recall()])
+ model.compile('sgd', loss='mse', metrics=[tf.keras.metrics.Recall()])
```
"""
@@ -1237,16 +1241,15 @@
dtype: (Optional) data type of the metric result.
"""
super(Recall, self).__init__(name=name, dtype=dtype)
- self.thresholds = 0.5 if thresholds is None else thresholds
- thresholds = to_list(thresholds)
- _assert_thresholds_range(thresholds)
+ self.thresholds = _parse_init_thresholds(
+ thresholds, default_threshold=0.5)
self.tp = self.add_weight(
'true_positives',
- shape=(len(thresholds),),
+ shape=(len(self.thresholds),),
initializer=init_ops.zeros_initializer)
self.fn = self.add_weight(
'false_negatives',
- shape=(len(thresholds),),
+ shape=(len(self.thresholds),),
initializer=init_ops.zeros_initializer)
def update_state(self, y_true, y_pred, sample_weight=None):
@@ -1269,7 +1272,7 @@
def result(self):
result = math_ops.div_no_nan(self.tp, self.tp + self.fn)
- return result if isinstance(self.thresholds, (list, tuple)) else result[0]
+ return result[0] if len(self.thresholds) == 1 else result
def reset_states(self):
num_thresholds = len(to_list(self.thresholds))
@@ -1341,6 +1344,7 @@
K.set_value(v, np.zeros((num_thresholds,)))
+@tf_export('keras.metrics.SensitivityAtSpecificity')
class SensitivityAtSpecificity(SensitivitySpecificityBase):
"""Computes the sensitivity at a given specificity.
@@ -1363,7 +1367,7 @@
Usage:
```python
- m = tf.metrics.SensitivityAtSpecificity(0.4, num_thresholds=1)
+ m = tf.keras.metrics.SensitivityAtSpecificity(0.4, num_thresholds=1)
m.update_state([0, 0, 1, 1], [0, 0.5, 0.3, 0.9])
print('Final result: ', m.result().numpy()) # Final result: 0.5
```
@@ -1375,7 +1379,7 @@
model.compile(
'sgd',
loss='mse',
- metrics=[tf.metrics.SensitivityAtSpecificity()])
+ metrics=[tf.keras.metrics.SensitivityAtSpecificity()])
```
"""
@@ -1409,6 +1413,7 @@
self.tp[min_index] + self.fn[min_index])
+@tf_export('keras.metrics.SpecificityAtSensitivity')
class SpecificityAtSensitivity(SensitivitySpecificityBase):
"""Computes the specificity at a given sensitivity.
@@ -1431,7 +1436,7 @@
Usage:
```python
- m = tf.metrics.SpecificityAtSensitivity(0.8, num_thresholds=1)
+ m = tf.keras.metrics.SpecificityAtSensitivity(0.8, num_thresholds=1)
m.update_state([0, 0, 1, 1], [0, 0.5, 0.3, 0.9])
print('Final result: ', m.result().numpy()) # Final result: 1.0
```
@@ -1443,7 +1448,7 @@
model.compile(
'sgd',
loss='mse',
- metrics=[tf.metrics.SpecificityAtSensitivity()])
+ metrics=[tf.keras.metrics.SpecificityAtSensitivity()])
```
"""
diff --git a/tensorflow/python/keras/metrics_functional_test.py b/tensorflow/python/keras/metrics_functional_test.py
new file mode 100644
index 0000000..513daaf
--- /dev/null
+++ b/tensorflow/python/keras/metrics_functional_test.py
@@ -0,0 +1,122 @@
+# 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.
+# ==============================================================================
+"""Tests for Keras metrics functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras import metrics
+from tensorflow.python.platform import test
+
+
+class KerasFunctionalMetricsTest(test.TestCase):
+
+ def test_metrics(self):
+ with self.cached_session():
+ y_a = K.variable(np.random.random((6, 7)))
+ y_b = K.variable(np.random.random((6, 7)))
+ for metric in [metrics.binary_accuracy, metrics.categorical_accuracy]:
+ output = metric(y_a, y_b)
+ self.assertEqual(K.eval(output).shape, (6,))
+
+ def test_sparse_categorical_accuracy_int(self):
+ with self.cached_session():
+ metric = metrics.sparse_categorical_accuracy
+ y_true = K.variable(np.random.randint(0, 7, (6,)))
+ y_pred = K.variable(np.random.random((6, 7)))
+ self.assertEqual(K.eval(metric(y_true, y_pred)).shape, (6,))
+
+ # Test correctness if the shape of y_true is (num_samples,)
+ y_true = K.variable([1., 0., 0., 0.])
+ y_pred = K.variable([[0.8, 0.2], [0.6, 0.4], [0.7, 0.3], [0.9, 0.1]])
+ print(K.eval(metric(y_true, y_pred)))
+ self.assertAllEqual(K.eval(metric(y_true, y_pred)), [0., 1., 1., 1.])
+
+ # Test correctness if the shape of y_true is (num_samples, 1)
+ y_true = K.variable([[1.], [0.], [0.], [0.]])
+ y_pred = K.variable([[0.8, 0.2], [0.6, 0.4], [0.7, 0.3], [0.9, 0.1]])
+ print(K.eval(metric(y_true, y_pred)))
+ self.assertAllEqual(K.eval(metric(y_true, y_pred)), [0., 1., 1., 1.])
+
+ def test_sparse_categorical_accuracy_float(self):
+ with self.cached_session():
+ metric = metrics.sparse_categorical_accuracy
+ y_true = K.variable(np.random.random((6,)))
+ y_pred = K.variable(np.random.random((6, 7)))
+ self.assertEqual(K.eval(metric(y_true, y_pred)).shape, (6,))
+
+ def test_sparse_categorical_accuracy_eager(self):
+ """Tests that ints passed in via Eager return results. See b/113504761."""
+ with context.eager_mode():
+ metric = metrics.sparse_categorical_accuracy
+ y_true = np.arange(6).reshape([6, 1])
+ y_pred = np.arange(36).reshape([6, 6])
+ self.assertAllEqual(metric(y_true, y_pred), [0., 0., 0., 0., 0., 1.])
+
+ def test_sparse_categorical_accuracy_float_eager(self):
+ """Tests that floats passed in via Eager return results. See b/113504761."""
+ with context.eager_mode():
+ metric = metrics.sparse_categorical_accuracy
+ y_true = np.arange(6, dtype=np.float32).reshape([6, 1])
+ y_pred = np.arange(36).reshape([6, 6])
+ self.assertAllEqual(metric(y_true, y_pred), [0., 0., 0., 0., 0., 1.])
+
+ def test_sparse_top_k_categorical_accuracy(self):
+ with self.cached_session():
+ # Test correctness if the shape of y_true is (num_samples, 1)
+ y_pred = K.variable(np.array([[0.3, 0.2, 0.1], [0.1, 0.2, 0.7]]))
+ y_true = K.variable(np.array([[1], [0]]))
+ result = K.eval(
+ metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=3))
+ self.assertEqual(result, 1)
+ result = K.eval(
+ metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=2))
+ self.assertEqual(result, 0.5)
+ result = K.eval(
+ metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=1))
+ self.assertEqual(result, 0.)
+
+ # Test correctness if the shape of y_true is (num_samples,)
+ y_pred = K.variable(np.array([[0.3, 0.2, 0.1], [0.1, 0.2, 0.7]]))
+ y_true = K.variable(np.array([1, 0]))
+ result = K.eval(
+ metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=3))
+ self.assertEqual(result, 1)
+ result = K.eval(
+ metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=2))
+ self.assertEqual(result, 0.5)
+ result = K.eval(
+ metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=1))
+ self.assertEqual(result, 0.)
+
+ def test_top_k_categorical_accuracy(self):
+ with self.cached_session():
+ y_pred = K.variable(np.array([[0.3, 0.2, 0.1], [0.1, 0.2, 0.7]]))
+ y_true = K.variable(np.array([[0, 1, 0], [1, 0, 0]]))
+ result = K.eval(metrics.top_k_categorical_accuracy(y_true, y_pred, k=3))
+ self.assertEqual(result, 1)
+ result = K.eval(metrics.top_k_categorical_accuracy(y_true, y_pred, k=2))
+ self.assertEqual(result, 0.5)
+ result = K.eval(metrics.top_k_categorical_accuracy(y_true, y_pred, k=1))
+ self.assertEqual(result, 0.)
+
+
+if __name__ == '__main__':
+ test.main()
diff --git a/tensorflow/python/keras/metrics_test.py b/tensorflow/python/keras/metrics_test.py
index 92398ac..9720d91 100644
--- a/tensorflow/python/keras/metrics_test.py
+++ b/tensorflow/python/keras/metrics_test.py
@@ -27,10 +27,10 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
-from tensorflow.python.keras import backend as K
+from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import layers
from tensorflow.python.keras import metrics
-from tensorflow.python.keras.models import Sequential
+from tensorflow.python.keras import testing_utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
@@ -40,98 +40,11 @@
from tensorflow.python.training.rmsprop import RMSPropOptimizer
-class KerasMetricsTest(test.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class KerasMeanTest(test.TestCase):
- def test_metrics(self):
- with self.cached_session():
- y_a = K.variable(np.random.random((6, 7)))
- y_b = K.variable(np.random.random((6, 7)))
- for metric in [metrics.binary_accuracy, metrics.categorical_accuracy]:
- output = metric(y_a, y_b)
- self.assertEqual(K.eval(output).shape, (6,))
-
- def test_sparse_categorical_accuracy_int(self):
- with self.cached_session():
- metric = metrics.sparse_categorical_accuracy
- y_true = K.variable(np.random.randint(0, 7, (6,)))
- y_pred = K.variable(np.random.random((6, 7)))
- self.assertEqual(K.eval(metric(y_true, y_pred)).shape, (6,))
-
- # Test correctness if the shape of y_true is (num_samples,)
- y_true = K.variable([1., 0., 0., 0.])
- y_pred = K.variable([[0.8, 0.2], [0.6, 0.4], [0.7, 0.3], [0.9, 0.1]])
- print(K.eval(metric(y_true, y_pred)))
- self.assertAllEqual(K.eval(metric(y_true, y_pred)), [0., 1., 1., 1.])
-
- # Test correctness if the shape of y_true is (num_samples, 1)
- y_true = K.variable([[1.], [0.], [0.], [0.]])
- y_pred = K.variable([[0.8, 0.2], [0.6, 0.4], [0.7, 0.3], [0.9, 0.1]])
- print(K.eval(metric(y_true, y_pred)))
- self.assertAllEqual(K.eval(metric(y_true, y_pred)), [0., 1., 1., 1.])
-
- def test_sparse_categorical_accuracy_float(self):
- with self.cached_session():
- metric = metrics.sparse_categorical_accuracy
- y_true = K.variable(np.random.random((6,)))
- y_pred = K.variable(np.random.random((6, 7)))
- self.assertEqual(K.eval(metric(y_true, y_pred)).shape, (6,))
-
- def test_sparse_categorical_accuracy_eager(self):
- """Tests that ints passed in via Eager return results. See b/113504761."""
- with context.eager_mode():
- metric = metrics.sparse_categorical_accuracy
- y_true = np.arange(6).reshape([6, 1])
- y_pred = np.arange(36).reshape([6, 6])
- self.assertAllEqual(metric(y_true, y_pred), [0., 0., 0., 0., 0., 1.])
-
- def test_sparse_categorical_accuracy_float_eager(self):
- """Tests that floats passed in via Eager return results. See b/113504761."""
- with context.eager_mode():
- metric = metrics.sparse_categorical_accuracy
- y_true = np.arange(6, dtype=np.float32).reshape([6, 1])
- y_pred = np.arange(36).reshape([6, 6])
- self.assertAllEqual(metric(y_true, y_pred), [0., 0., 0., 0., 0., 1.])
-
- def test_sparse_top_k_categorical_accuracy(self):
- with self.cached_session():
- # Test correctness if the shape of y_true is (num_samples, 1)
- y_pred = K.variable(np.array([[0.3, 0.2, 0.1], [0.1, 0.2, 0.7]]))
- y_true = K.variable(np.array([[1], [0]]))
- result = K.eval(
- metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=3))
- self.assertEqual(result, 1)
- result = K.eval(
- metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=2))
- self.assertEqual(result, 0.5)
- result = K.eval(
- metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=1))
- self.assertEqual(result, 0.)
-
- # Test correctness if the shape of y_true is (num_samples,)
- y_pred = K.variable(np.array([[0.3, 0.2, 0.1], [0.1, 0.2, 0.7]]))
- y_true = K.variable(np.array([1, 0]))
- result = K.eval(
- metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=3))
- self.assertEqual(result, 1)
- result = K.eval(
- metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=2))
- self.assertEqual(result, 0.5)
- result = K.eval(
- metrics.sparse_top_k_categorical_accuracy(y_true, y_pred, k=1))
- self.assertEqual(result, 0.)
-
- def test_top_k_categorical_accuracy(self):
- with self.cached_session():
- y_pred = K.variable(np.array([[0.3, 0.2, 0.1], [0.1, 0.2, 0.7]]))
- y_true = K.variable(np.array([[0, 1, 0], [1, 0, 0]]))
- result = K.eval(metrics.top_k_categorical_accuracy(y_true, y_pred, k=3))
- self.assertEqual(result, 1)
- result = K.eval(metrics.top_k_categorical_accuracy(y_true, y_pred, k=2))
- self.assertEqual(result, 0.5)
- result = K.eval(metrics.top_k_categorical_accuracy(y_true, y_pred, k=1))
- self.assertEqual(result, 0.)
-
- @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True)
+ # TODO(b/120949004): Re-enable garbage collection check
+ # @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True)
def test_mean(self):
m = metrics.Mean(name='my_mean')
@@ -163,7 +76,6 @@
self.assertEqual(self.evaluate(m.total), 0)
self.assertEqual(self.evaluate(m.count), 0)
- @test_util.run_in_graph_and_eager_modes
def test_mean_with_sample_weight(self):
m = metrics.Mean(dtype=dtypes.float64)
self.assertEqual(m.dtype, dtypes.float64)
@@ -227,7 +139,6 @@
self.assertAlmostEqual(self.evaluate(m.count), 1.7, 2) # 0.5 + 1.2
self.assertAlmostEqual(result, 52 / 1.7, 2)
- @test_util.run_in_graph_and_eager_modes
def test_save_restore(self):
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt')
@@ -258,7 +169,10 @@
self.assertEqual(200., self.evaluate(restore_mean.result()))
self.assertEqual(3, self.evaluate(restore_mean.count))
- @test_util.run_in_graph_and_eager_modes
+
+@test_util.run_all_in_graph_and_eager_modes
+class KerasAccuracyTest(test.TestCase):
+
def test_accuracy(self):
acc_obj = metrics.Accuracy(name='my acc')
@@ -280,7 +194,6 @@
result = self.evaluate(result_t)
self.assertAlmostEqual(result, 0.96, 2) # 4.5/4.7
- @test_util.run_in_graph_and_eager_modes
def test_binary_accuracy(self):
acc_obj = metrics.BinaryAccuracy(name='my acc')
@@ -313,7 +226,6 @@
result = self.evaluate(result_t)
self.assertAlmostEqual(result, 0.67, 2) # 4.5/6.7
- @test_util.run_in_graph_and_eager_modes
def test_binary_accuracy_threshold(self):
acc_obj = metrics.BinaryAccuracy(threshold=0.7)
self.evaluate(variables.variables_initializer(acc_obj.variables))
@@ -321,7 +233,6 @@
result = self.evaluate(result_t)
self.assertAlmostEqual(result, 0.5, 2)
- @test_util.run_in_graph_and_eager_modes
def test_categorical_accuracy(self):
acc_obj = metrics.CategoricalAccuracy(name='my acc')
@@ -345,7 +256,6 @@
result = self.evaluate(result_t)
self.assertAlmostEqual(result, 0.93, 2) # 2.5/2.7
- @test_util.run_in_graph_and_eager_modes
def test_sparse_categorical_accuracy(self):
acc_obj = metrics.SparseCategoricalAccuracy(name='my acc')
@@ -369,18 +279,11 @@
result = self.evaluate(result_t)
self.assertAlmostEqual(result, 0.93, 2) # 2.5/2.7
-
-def _get_simple_sequential_model(compile_metrics):
- model = Sequential()
- model.add(
- layers.Dense(
- 3, activation='relu', input_dim=4, kernel_initializer='ones'))
- model.add(layers.Dense(1, activation='sigmoid', kernel_initializer='ones'))
- model.compile(
- loss='mae',
- metrics=compile_metrics,
- optimizer=RMSPropOptimizer(learning_rate=0.001))
- return model
+ def test_assert_thresholds_range(self):
+ with self.assertRaisesRegexp(
+ ValueError,
+ r'Threshold values must be in \[0, 1\]. Invalid values: \[None\]'):
+ metrics._assert_thresholds_range([None, 0.5])
@test_util.run_all_in_graph_and_eager_modes
@@ -451,16 +354,6 @@
r'Threshold values must be in \[0, 1\]. Invalid values: \[-1, 2\]'):
metrics.FalsePositives(thresholds=[-1, 0.5, 2])
- def test_reset_states(self):
- fp_obj = metrics.FalsePositives()
- model = _get_simple_sequential_model([fp_obj])
- x = np.ones((100, 4))
- y = np.zeros((100, 1))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(fp_obj.accumulator), 100.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(fp_obj.accumulator), 100.)
-
@test_util.run_all_in_graph_and_eager_modes
class FalseNegativesTest(test.TestCase):
@@ -523,16 +416,6 @@
result = fn_obj(y_true, y_pred, sample_weight=sample_weight)
self.assertAllClose([4., 16., 23.], self.evaluate(result))
- def test_reset_states(self):
- fn_obj = metrics.FalseNegatives()
- model = _get_simple_sequential_model([fn_obj])
- x = np.zeros((100, 4))
- y = np.ones((100, 1))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(fn_obj.accumulator), 100.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(fn_obj.accumulator), 100.)
-
@test_util.run_all_in_graph_and_eager_modes
class TrueNegativesTest(test.TestCase):
@@ -595,16 +478,6 @@
result = tn_obj(y_true, y_pred, sample_weight=sample_weight)
self.assertAllClose([5., 15., 23.], self.evaluate(result))
- def test_reset_states(self):
- tn_obj = metrics.TrueNegatives()
- model = _get_simple_sequential_model([tn_obj])
- x = np.zeros((100, 4))
- y = np.zeros((100, 1))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(tn_obj.accumulator), 100.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(tn_obj.accumulator), 100.)
-
@test_util.run_all_in_graph_and_eager_modes
class TruePositivesTest(test.TestCase):
@@ -666,16 +539,6 @@
result = tp_obj(y_true, y_pred, sample_weight=37.)
self.assertAllClose([222., 111., 37.], self.evaluate(result))
- def test_reset_states(self):
- tp_obj = metrics.TruePositives()
- model = _get_simple_sequential_model([tp_obj])
- x = np.ones((100, 4))
- y = np.ones((100, 1))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(tp_obj.accumulator), 100.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(tp_obj.accumulator), 100.)
-
@test_util.run_all_in_graph_and_eager_modes
class PrecisionTest(test.TestCase):
@@ -683,7 +546,7 @@
def test_config(self):
p_obj = metrics.Precision(name='my_precision', thresholds=[0.4, 0.9])
self.assertEqual(p_obj.name, 'my_precision')
- self.assertLen(p_obj.variables, 2)
+ self.assertEqual(len(p_obj.variables), 2)
self.assertEqual([v.name for v in p_obj.variables],
['true_positives:0', 'false_positives:0'])
self.assertEqual(p_obj.thresholds, [0.4, 0.9])
@@ -788,18 +651,6 @@
self.assertArrayNear([expected_precision, 0], self.evaluate(p_obj.result()),
1e-3)
- def test_reset_states(self):
- p_obj = metrics.Precision()
- model = _get_simple_sequential_model([p_obj])
- x = np.concatenate((np.ones((50, 4)), np.ones((50, 4))))
- y = np.concatenate((np.ones((50, 1)), np.zeros((50, 1))))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(p_obj.tp), 50.)
- self.assertEqual(self.evaluate(p_obj.fp), 50.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(p_obj.tp), 50.)
- self.assertEqual(self.evaluate(p_obj.fp), 50.)
-
@test_util.run_all_in_graph_and_eager_modes
class RecallTest(test.TestCase):
@@ -807,7 +658,7 @@
def test_config(self):
r_obj = metrics.Recall(name='my_recall', thresholds=[0.4, 0.9])
self.assertEqual(r_obj.name, 'my_recall')
- self.assertLen(r_obj.variables, 2)
+ self.assertEqual(len(r_obj.variables), 2)
self.assertEqual([v.name for v in r_obj.variables],
['true_positives:0', 'false_negatives:0'])
self.assertEqual(r_obj.thresholds, [0.4, 0.9])
@@ -911,18 +762,6 @@
self.assertArrayNear([expected_recall, 0], self.evaluate(r_obj.result()),
1e-3)
- def test_reset_states(self):
- r_obj = metrics.Recall()
- model = _get_simple_sequential_model([r_obj])
- x = np.concatenate((np.ones((50, 4)), np.zeros((50, 4))))
- y = np.concatenate((np.ones((50, 1)), np.ones((50, 1))))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(r_obj.tp), 50.)
- self.assertEqual(self.evaluate(r_obj.fn), 50.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(r_obj.tp), 50.)
- self.assertEqual(self.evaluate(r_obj.fn), 50.)
-
@test_util.run_all_in_graph_and_eager_modes
class SensitivityAtSpecificityTest(test.TestCase, parameterized.TestCase):
@@ -1012,24 +851,6 @@
with self.assertRaisesRegexp(ValueError, '`num_thresholds` must be > 0.'):
metrics.SensitivityAtSpecificity(0.4, num_thresholds=-1)
- def test_reset_states(self):
- s_obj = metrics.SensitivityAtSpecificity(0.5, num_thresholds=1)
- model = _get_simple_sequential_model([s_obj])
- x = np.concatenate((np.ones((25, 4)), np.zeros((25, 4)), np.zeros((25, 4)),
- np.ones((25, 4))))
- y = np.concatenate((np.ones((25, 1)), np.zeros((25, 1)), np.ones((25, 1)),
- np.zeros((25, 1))))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(s_obj.tp), 25.)
- self.assertEqual(self.evaluate(s_obj.fp), 25.)
- self.assertEqual(self.evaluate(s_obj.fn), 25.)
- self.assertEqual(self.evaluate(s_obj.tn), 25.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(s_obj.tp), 25.)
- self.assertEqual(self.evaluate(s_obj.fp), 25.)
- self.assertEqual(self.evaluate(s_obj.fn), 25.)
- self.assertEqual(self.evaluate(s_obj.tn), 25.)
-
@test_util.run_all_in_graph_and_eager_modes
class SpecificityAtSensitivityTest(test.TestCase, parameterized.TestCase):
@@ -1119,24 +940,6 @@
with self.assertRaisesRegexp(ValueError, '`num_thresholds` must be > 0.'):
metrics.SpecificityAtSensitivity(0.4, num_thresholds=-1)
- def test_reset_states(self):
- s_obj = metrics.SpecificityAtSensitivity(0.5, num_thresholds=1)
- model = _get_simple_sequential_model([s_obj])
- x = np.concatenate((np.ones((25, 4)), np.zeros((25, 4)), np.zeros((25, 4)),
- np.ones((25, 4))))
- y = np.concatenate((np.ones((25, 1)), np.zeros((25, 1)), np.ones((25, 1)),
- np.zeros((25, 1))))
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(s_obj.tp), 25.)
- self.assertEqual(self.evaluate(s_obj.fp), 25.)
- self.assertEqual(self.evaluate(s_obj.fn), 25.)
- self.assertEqual(self.evaluate(s_obj.tn), 25.)
- model.evaluate(x, y)
- self.assertEqual(self.evaluate(s_obj.tp), 25.)
- self.assertEqual(self.evaluate(s_obj.fp), 25.)
- self.assertEqual(self.evaluate(s_obj.fn), 25.)
- self.assertEqual(self.evaluate(s_obj.tn), 25.)
-
@test_util.run_all_in_graph_and_eager_modes
class CosineProximityTest(test.TestCase):
@@ -1171,5 +974,125 @@
result = cosine_obj(y_true, y_pred, sample_weight=sample_weight)
self.assertAllClose(-0.59916, self.evaluate(result), atol=1e-5)
+
+def _get_model(compile_metrics):
+ model_layers = [
+ layers.Dense(3, activation='relu', kernel_initializer='ones'),
+ layers.Dense(1, activation='sigmoid', kernel_initializer='ones')]
+
+ model = testing_utils.get_model_from_layers(model_layers, input_shape=(4,))
+ model.compile(
+ loss='mae',
+ metrics=compile_metrics,
+ optimizer=RMSPropOptimizer(learning_rate=0.001),
+ run_eagerly=testing_utils.should_run_eagerly())
+ return model
+
+
+@keras_parameterized.run_with_all_model_types
+@keras_parameterized.run_all_keras_modes
+class ResetStatesTest(keras_parameterized.TestCase):
+
+ def test_reset_states_false_positives(self):
+ fp_obj = metrics.FalsePositives()
+ model = _get_model([fp_obj])
+ x = np.ones((100, 4))
+ y = np.zeros((100, 1))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(fp_obj.accumulator), 100.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(fp_obj.accumulator), 100.)
+
+ def test_reset_states_false_negatives(self):
+ fn_obj = metrics.FalseNegatives()
+ model = _get_model([fn_obj])
+ x = np.zeros((100, 4))
+ y = np.ones((100, 1))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(fn_obj.accumulator), 100.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(fn_obj.accumulator), 100.)
+
+ def test_reset_states_true_negatives(self):
+ tn_obj = metrics.TrueNegatives()
+ model = _get_model([tn_obj])
+ x = np.zeros((100, 4))
+ y = np.zeros((100, 1))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(tn_obj.accumulator), 100.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(tn_obj.accumulator), 100.)
+
+ def test_reset_states_true_positives(self):
+ tp_obj = metrics.TruePositives()
+ model = _get_model([tp_obj])
+ x = np.ones((100, 4))
+ y = np.ones((100, 1))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(tp_obj.accumulator), 100.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(tp_obj.accumulator), 100.)
+
+ def test_reset_states_precision(self):
+ p_obj = metrics.Precision()
+ model = _get_model([p_obj])
+ x = np.concatenate((np.ones((50, 4)), np.ones((50, 4))))
+ y = np.concatenate((np.ones((50, 1)), np.zeros((50, 1))))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(p_obj.tp), 50.)
+ self.assertEqual(self.evaluate(p_obj.fp), 50.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(p_obj.tp), 50.)
+ self.assertEqual(self.evaluate(p_obj.fp), 50.)
+
+ def test_reset_states_recall(self):
+ r_obj = metrics.Recall()
+ model = _get_model([r_obj])
+ x = np.concatenate((np.ones((50, 4)), np.zeros((50, 4))))
+ y = np.concatenate((np.ones((50, 1)), np.ones((50, 1))))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(r_obj.tp), 50.)
+ self.assertEqual(self.evaluate(r_obj.fn), 50.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(r_obj.tp), 50.)
+ self.assertEqual(self.evaluate(r_obj.fn), 50.)
+
+ def test_reset_states_sensitivity_at_specificity(self):
+ s_obj = metrics.SensitivityAtSpecificity(0.5, num_thresholds=1)
+ model = _get_model([s_obj])
+ x = np.concatenate((np.ones((25, 4)), np.zeros((25, 4)), np.zeros((25, 4)),
+ np.ones((25, 4))))
+ y = np.concatenate((np.ones((25, 1)), np.zeros((25, 1)), np.ones((25, 1)),
+ np.zeros((25, 1))))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(s_obj.tp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fn), 25.)
+ self.assertEqual(self.evaluate(s_obj.tn), 25.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(s_obj.tp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fn), 25.)
+ self.assertEqual(self.evaluate(s_obj.tn), 25.)
+
+ def test_reset_states_specificity_at_sensitivity(self):
+ s_obj = metrics.SpecificityAtSensitivity(0.5, num_thresholds=1)
+ model = _get_model([s_obj])
+ x = np.concatenate((np.ones((25, 4)), np.zeros((25, 4)), np.zeros((25, 4)),
+ np.ones((25, 4))))
+ y = np.concatenate((np.ones((25, 1)), np.zeros((25, 1)), np.ones((25, 1)),
+ np.zeros((25, 1))))
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(s_obj.tp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fn), 25.)
+ self.assertEqual(self.evaluate(s_obj.tn), 25.)
+ model.evaluate(x, y)
+ self.assertEqual(self.evaluate(s_obj.tp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fp), 25.)
+ self.assertEqual(self.evaluate(s_obj.fn), 25.)
+ self.assertEqual(self.evaluate(s_obj.tn), 25.)
+
+
if __name__ == '__main__':
test.main()
diff --git a/tensorflow/python/keras/model_subclassing_test.py b/tensorflow/python/keras/model_subclassing_test.py
index 620275e..553c7fb 100644
--- a/tensorflow/python/keras/model_subclassing_test.py
+++ b/tensorflow/python/keras/model_subclassing_test.py
@@ -187,6 +187,7 @@
@test_util.run_all_in_graph_and_eager_modes
+@test_util.run_v1_only('b/120545219')
class ModelSubclassingTest(test.TestCase):
def test_custom_build(self):
@@ -915,6 +916,7 @@
self.assertEqual(1, len(model.get_updates_for(x)))
+@test_util.run_v1_only('b/120545219')
class GraphSpecificModelSubclassingTests(test.TestCase):
@test_util.run_deprecated_v1
diff --git a/tensorflow/python/keras/models_test.py b/tensorflow/python/keras/models_test.py
index b0872ae..fe7d8a5 100644
--- a/tensorflow/python/keras/models_test.py
+++ b/tensorflow/python/keras/models_test.py
@@ -31,6 +31,7 @@
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import metrics
from tensorflow.python.keras import models
+from tensorflow.python.keras import optimizers
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import resource_variable_ops
@@ -69,7 +70,7 @@
class TestModelCloning(test.TestCase):
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_clone_sequential_model(self):
with self.cached_session():
val_a = np.random.random((10, 4))
@@ -102,10 +103,9 @@
new_model.compile('rmsprop', 'mse')
new_model.train_on_batch(None, val_out)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_clone_sequential_model_input_layer(self):
- @test_util.run_deprecated_v1
def test_input_layer(include_inputs):
with self.cached_session():
val_a = np.random.random((10, 4))
@@ -142,7 +142,7 @@
test_input_layer(True)
test_input_layer(False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def test_clone_functional_model(self):
with self.cached_session():
val_a = np.random.random((10, 4))
@@ -318,6 +318,7 @@
model_copy.get_weights()[0]))
+@test_util.run_v1_only('b/120545219')
class TestCloneAndBuildModel(test.TestCase):
def test_clone_and_build_non_compiled_model(self):
@@ -365,7 +366,9 @@
self.assertEqual('mse', model.loss)
self.assertTrue(
- isinstance(model.optimizer, keras.optimizers.RMSprop))
+ isinstance(model.optimizer,
+ (keras.optimizers.RMSprop,
+ keras.optimizer_v2.rmsprop.RMSprop)))
self.assertEqual(['acc', metrics.categorical_accuracy],
model._compile_metrics)
@@ -402,7 +405,6 @@
new_model.train_on_batch(inp, out)
new_model.evaluate(inp, out)
- @test_util.run_deprecated_v1
def test_clone_and_build_compiled_sequential_model(self):
with self.cached_session():
model = keras.models.Sequential()
@@ -415,7 +417,6 @@
self._clone_and_build_test_helper(model)
- @test_util.run_deprecated_v1
def test_clone_and_build_functional_model(self):
with self.cached_session():
input_a = keras.Input(shape=(4,))
@@ -432,7 +433,6 @@
self._clone_and_build_test_helper(model)
- @test_util.run_deprecated_v1
def test_clone_and_build_subclassed_model(self):
class SubclassedModel(keras.Model):
@@ -481,11 +481,9 @@
def test_replace_tf_optimizer_iterations_variable(self):
self.assert_optimizer_iterations_increases(adam.AdamOptimizer(0.01))
- @test_util.run_deprecated_v1
def test_replace_keras_optimizer_iterations_variable(self):
- self.assert_optimizer_iterations_increases('adam')
+ self.assert_optimizer_iterations_increases(optimizers.Adam())
- @test_util.run_deprecated_v1
def test_clone_and_build_sequential_model_without_inputs_defined(self):
with self.cached_session():
model = sequential_model(False, False)
diff --git a/tensorflow/python/keras/optimizer_v2/adadelta.py b/tensorflow/python/keras/optimizer_v2/adadelta.py
index 55b4eba..8985325 100644
--- a/tensorflow/python/keras/optimizer_v2/adadelta.py
+++ b/tensorflow/python/keras/optimizer_v2/adadelta.py
@@ -22,8 +22,10 @@
from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export('keras.optimizers.Adadelta')
class Adadelta(optimizer_v2.OptimizerV2):
r"""Optimizer that implements the Adadelta algorithm.
@@ -85,7 +87,7 @@
@end_compatibility
"""
super(Adadelta, self).__init__(name, **kwargs)
- self._set_hyper('learning_rate', learning_rate)
+ self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
self._set_hyper('decay', self._initial_decay)
self._set_hyper('rho', rho)
self._set_hyper('epsilon', epsilon)
diff --git a/tensorflow/python/keras/optimizer_v2/adadelta_test.py b/tensorflow/python/keras/optimizer_v2/adadelta_test.py
index 0fb67d0..c95af6a 100644
--- a/tensorflow/python/keras/optimizer_v2/adadelta_test.py
+++ b/tensorflow/python/keras/optimizer_v2/adadelta_test.py
@@ -153,8 +153,11 @@
with self.cached_session():
var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
- pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
- loss = pred * pred
+
+ def loss():
+ pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x) # pylint: disable=cell-var-from-loop
+ return pred * pred
+
sgd_op = adadelta.Adadelta(1.0, 1.0, 1.0).minimize(
loss, var_list=[var0])
variables.global_variables_initializer().run()
@@ -165,6 +168,14 @@
# Validate updated params
self.assertAllCloseAccordingToType([[-111, -138]], self.evaluate(var0))
+ def testConstructAdadeltaWithLR(self):
+ opt = adadelta.Adadelta(lr=1.0, rho=0.9, epsilon=1.)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = adadelta.Adadelta(learning_rate=0.1, rho=0.9, epsilon=1., lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = adadelta.Adadelta(learning_rate=0.1, rho=0.9, epsilon=1.)
+ self.assertEqual(opt_3.lr, 0.1)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizer_v2/adagrad.py b/tensorflow/python/keras/optimizer_v2/adagrad.py
index 670cad7..6faf7fc 100644
--- a/tensorflow/python/keras/optimizer_v2/adagrad.py
+++ b/tensorflow/python/keras/optimizer_v2/adagrad.py
@@ -27,8 +27,10 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import state_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export('keras.optimizers.Adagrad')
class Adagrad(optimizer_v2.OptimizerV2):
r"""Optimizer that implements the Adagrad algorithm.
@@ -86,7 +88,7 @@
if epsilon < 1e-7:
raise ValueError('epsilon must be larger than 1e-7: %s' % epsilon)
super(Adagrad, self).__init__(name, **kwargs)
- self._set_hyper('learning_rate', learning_rate)
+ self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
self._set_hyper('decay', self._initial_decay)
self._initial_accumulator_value = initial_accumulator_value
self._set_hyper('epsilon', epsilon)
diff --git a/tensorflow/python/keras/optimizer_v2/adagrad_test.py b/tensorflow/python/keras/optimizer_v2/adagrad_test.py
index b2c2901..cf6f6a7 100644
--- a/tensorflow/python/keras/optimizer_v2/adagrad_test.py
+++ b/tensorflow/python/keras/optimizer_v2/adagrad_test.py
@@ -167,8 +167,11 @@
var0 = resource_variable_ops.ResourceVariable(
[[1.0, 2.0], [3.0, 4.0]], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
- pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
- loss = pred * pred
+
+ def loss():
+ pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x) # pylint: disable=cell-var-from-loop
+ return pred * pred
+
sgd_op = adagrad.Adagrad(1.0).minimize(loss, var_list=[var0])
variables.global_variables_initializer().run()
# Fetch params to validate initial values
@@ -297,12 +300,12 @@
with self.cached_session():
var_repeated = resource_variable_ops.ResourceVariable(
[1.0, 2.0], dtype=dtype)
- loss_repeated = math_ops.reduce_sum(
- embedding_ops.embedding_lookup(var_repeated, [0, 0]))
+ loss_repeated = lambda: math_ops.reduce_sum( # pylint: disable=g-long-lambda
+ embedding_ops.embedding_lookup(var_repeated, [0, 0])) # pylint: disable=cell-var-from-loop
var_aggregated = resource_variable_ops.ResourceVariable(
[1.0, 2.0], dtype=dtype)
- loss_aggregated = 2 * math_ops.reduce_sum(
- embedding_ops.embedding_lookup(var_aggregated, [0]))
+ loss_aggregated = lambda: 2 * math_ops.reduce_sum( # pylint: disable=g-long-lambda
+ embedding_ops.embedding_lookup(var_aggregated, [0])) # pylint: disable=cell-var-from-loop
update_op_repeated = adagrad.Adagrad(2.0).minimize(
loss_repeated, var_list=[var_repeated])
update_op_aggregated = adagrad.Adagrad(2.0).minimize(
@@ -395,6 +398,14 @@
self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+ def testConstructAdagradWithLR(self):
+ opt = adagrad.Adagrad(lr=1.0)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = adagrad.Adagrad(learning_rate=0.1, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = adagrad.Adagrad(learning_rate=0.1)
+ self.assertEqual(opt_3.lr, 0.1)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizer_v2/adam.py b/tensorflow/python/keras/optimizer_v2/adam.py
index ef3d783..c99468f 100644
--- a/tensorflow/python/keras/optimizer_v2/adam.py
+++ b/tensorflow/python/keras/optimizer_v2/adam.py
@@ -24,8 +24,10 @@
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export('keras.optimizers.Adam')
class Adam(optimizer_v2.OptimizerV2):
"""Optimizer that implements the Adam algorithm.
@@ -127,12 +129,12 @@
"""
super(Adam, self).__init__(name, **kwargs)
- self._set_hyper('learning_rate', learning_rate)
+ self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
self._set_hyper('decay', self._initial_decay)
self._set_hyper('beta_1', beta_1)
self._set_hyper('beta_2', beta_2)
self._set_hyper('epsilon', epsilon)
- self._amsgrad = amsgrad
+ self.amsgrad = amsgrad
def _create_slots(self, var_list):
# Create slots for the first and second moments.
@@ -141,7 +143,7 @@
self.add_slot(var, 'm')
for var in var_list:
self.add_slot(var, 'v')
- if self._amsgrad:
+ if self.amsgrad:
for var in var_list:
self.add_slot(var, 'vhat')
@@ -166,7 +168,7 @@
local_step = math_ops.cast(self.iterations + 1, var_dtype)
beta_1_power = math_ops.pow(beta_1_t, local_step)
beta_2_power = math_ops.pow(beta_2_t, local_step)
- if not self._amsgrad:
+ if not self.amsgrad:
return training_ops.resource_apply_adam(
var.handle,
m.handle,
@@ -220,7 +222,7 @@
with ops.control_dependencies([v_t]):
v_t = self._resource_scatter_add(v, indices, v_scaled_g_values)
- if not self._amsgrad:
+ if not self.amsgrad:
v_sqrt = math_ops.sqrt(v_t)
var_update = state_ops.assign_sub(
var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
@@ -251,6 +253,6 @@
'beta_1': self._serialize_hyperparameter('beta_1'),
'beta_2': self._serialize_hyperparameter('beta_2'),
'epsilon': self._serialize_hyperparameter('epsilon'),
- 'amsgrad': self._amsgrad,
+ 'amsgrad': self.amsgrad,
})
return config
diff --git a/tensorflow/python/keras/optimizer_v2/adam_test.py b/tensorflow/python/keras/optimizer_v2/adam_test.py
index 3bbafe1..49a9de4 100644
--- a/tensorflow/python/keras/optimizer_v2/adam_test.py
+++ b/tensorflow/python/keras/optimizer_v2/adam_test.py
@@ -162,9 +162,9 @@
# it (i.e. they have GPU kernels).
var = variables.Variable([[1.0], [2.0]])
indices = constant_op.constant([0, 1], dtype=index_dtype)
- gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices))
+ g_sum = lambda: math_ops.reduce_sum(array_ops.gather(var, indices)) # pylint: disable=cell-var-from-loop
optimizer = adam.Adam(3.0)
- minimize_op = optimizer.minimize(gathered_sum, var_list=[var])
+ minimize_op = optimizer.minimize(g_sum, var_list=[var])
variables.global_variables_initializer().run()
minimize_op.run()
@@ -503,6 +503,14 @@
self.assertEqual(
self.evaluate(keras_v1_iteration), self.evaluate(keras_v2_iteration))
+ def testConstructAdamWithLR(self):
+ opt = adam.Adam(lr=1.0)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = adam.Adam(learning_rate=0.1, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = adam.Adam(learning_rate=0.1)
+ self.assertEqual(opt_3.lr, 0.1)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizer_v2/adamax.py b/tensorflow/python/keras/optimizer_v2/adamax.py
index ddd7858..920a6c0 100644
--- a/tensorflow/python/keras/optimizer_v2/adamax.py
+++ b/tensorflow/python/keras/optimizer_v2/adamax.py
@@ -25,8 +25,10 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export('keras.optimizers.Adamax')
class Adamax(adam.Adam):
"""Optimizer that implements the Adamax algorithm.
diff --git a/tensorflow/python/keras/optimizer_v2/adamax_test.py b/tensorflow/python/keras/optimizer_v2/adamax_test.py
index baf131f..339c0fe 100644
--- a/tensorflow/python/keras/optimizer_v2/adamax_test.py
+++ b/tensorflow/python/keras/optimizer_v2/adamax_test.py
@@ -136,9 +136,9 @@
# it (i.e. they have GPU kernels).
var = variables.Variable([[1.0], [2.0]])
indices = constant_op.constant([0, 1], dtype=index_dtype)
- gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices))
+ g_sum = lambda: math_ops.reduce_sum(array_ops.gather(var, indices)) # pylint: disable=cell-var-from-loop
optimizer = adamax.Adamax(3.0)
- minimize_op = optimizer.minimize(gathered_sum, var_list=[var])
+ minimize_op = optimizer.minimize(g_sum, var_list=[var])
variables.global_variables_initializer().run()
minimize_op.run()
@@ -362,6 +362,14 @@
# There should be iteration, and two unique slot variables for v1 and v2.
self.assertEqual(5, len(set(opt.variables())))
+ def testConstructAdamaxWithLR(self):
+ opt = adamax.Adamax(lr=1.0)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = adamax.Adamax(learning_rate=0.1, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = adamax.Adamax(learning_rate=0.1)
+ self.assertEqual(opt_3.lr, 0.1)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizer_v2/ftrl.py b/tensorflow/python/keras/optimizer_v2/ftrl.py
index e278e35..365bd68 100644
--- a/tensorflow/python/keras/optimizer_v2/ftrl.py
+++ b/tensorflow/python/keras/optimizer_v2/ftrl.py
@@ -21,8 +21,10 @@
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export('keras.optimizers.Ftrl')
class Ftrl(optimizer_v2.OptimizerV2):
"""Optimizer that implements the FTRL algorithm.
diff --git a/tensorflow/python/keras/optimizer_v2/ftrl_test.py b/tensorflow/python/keras/optimizer_v2/ftrl_test.py
index bec400e..f0f07e9 100644
--- a/tensorflow/python/keras/optimizer_v2/ftrl_test.py
+++ b/tensorflow/python/keras/optimizer_v2/ftrl_test.py
@@ -113,8 +113,11 @@
with self.cached_session():
var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
- pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
- loss = pred * pred
+
+ def loss():
+ pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x) # pylint: disable=cell-var-from-loop
+ return pred * pred
+
sgd_op = ftrl.Ftrl(1.0).minimize(loss, var_list=[var0])
variables.global_variables_initializer().run()
# Fetch params to validate initial values
diff --git a/tensorflow/python/keras/optimizer_v2/gradient_descent.py b/tensorflow/python/keras/optimizer_v2/gradient_descent.py
index 2b82b5e..a77ae30 100644
--- a/tensorflow/python/keras/optimizer_v2/gradient_descent.py
+++ b/tensorflow/python/keras/optimizer_v2/gradient_descent.py
@@ -1,4 +1,4 @@
-# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -21,8 +21,10 @@
from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export("keras.optimizers.SGD")
class SGD(optimizer_v2.OptimizerV2):
"""Stochastic gradient descent and momentum optimizer.
@@ -32,7 +34,7 @@
gradient is evaluated at theta(t).
```
- or Computes (if `use_nesterov = False`):
+ or Computes (if `nesterov = False`):
```
v(t+1) = momentum * v(t) - learning_rate * gradient
theta(t+1) = theta(t) + v(t+1)
@@ -75,7 +77,7 @@
**kwargs: keyword arguments. Allowed to be {`decay`}
"""
super(SGD, self).__init__(name, **kwargs)
- self._set_hyper("learning_rate", learning_rate)
+ self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
self._set_hyper("decay", self._initial_decay)
self._momentum = False
@@ -85,7 +87,7 @@
raise ValueError("`momentum` must be between [0, 1].")
self._set_hyper("momentum", momentum)
- self._nesterov = nesterov
+ self.nesterov = nesterov
def _create_slots(self, var_list):
if self._momentum:
@@ -104,7 +106,7 @@
grad,
self._get_hyper("momentum", var_dtype),
use_locking=self._use_locking,
- use_nesterov=self._nesterov)
+ use_nesterov=self.nesterov)
else:
return training_ops.resource_apply_gradient_descent(
var.handle, lr_t, grad, use_locking=self._use_locking)
@@ -132,7 +134,7 @@
indices,
self._get_hyper("momentum", var_dtype),
use_locking=self._use_locking,
- use_nesterov=self._nesterov)
+ use_nesterov=self.nesterov)
def get_config(self):
config = super(SGD, self).get_config()
@@ -140,6 +142,6 @@
"learning_rate": self._serialize_hyperparameter("learning_rate"),
"decay": self._serialize_hyperparameter("decay"),
"momentum": self._serialize_hyperparameter("momentum"),
- "nesterov": self._nesterov,
+ "nesterov": self.nesterov,
})
return config
diff --git a/tensorflow/python/keras/optimizer_v2/gradient_descent_test.py b/tensorflow/python/keras/optimizer_v2/gradient_descent_test.py
index 0c64202..9a4178d 100644
--- a/tensorflow/python/keras/optimizer_v2/gradient_descent_test.py
+++ b/tensorflow/python/keras/optimizer_v2/gradient_descent_test.py
@@ -122,8 +122,6 @@
var1 = resource_variable_ops.ResourceVariable([3.0], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
loss = lambda: math_ops.matmul(var0, x) + var1 # pylint: disable=cell-var-from-loop
- if not context.executing_eagerly():
- loss = loss()
sgd = gradient_descent.SGD(1.0)
sgd_op = sgd.minimize(loss, [var0, var1])
self.evaluate(variables.global_variables_initializer())
@@ -141,9 +139,12 @@
var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
var1 = resource_variable_ops.ResourceVariable([3.0], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
- pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
- pred += var1
- loss = pred * pred
+
+ def loss():
+ pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x) # pylint: disable=cell-var-from-loop
+ pred += var1 # pylint: disable=cell-var-from-loop
+ return pred * pred
+
sgd_op = gradient_descent.SGD(1.0).minimize(loss, [var0, var1])
self.evaluate(variables.global_variables_initializer())
# Run 1 step of sgd
@@ -181,7 +182,8 @@
opt = gradient_descent.SGD(3.0)
values = [1.0, 3.0]
vars_ = [variables.Variable([v], dtype=dtype) for v in values]
- grads_and_vars = opt.compute_gradients(vars_[0] + vars_[1], vars_)
+ loss = lambda: vars_[0] + vars_[1] # pylint: disable=cell-var-from-loop
+ grads_and_vars = opt._compute_gradients(loss, vars_)
self.evaluate(variables.global_variables_initializer())
for grad, _ in grads_and_vars:
self.assertAllCloseAccordingToType([1.0], self.evaluate(grad))
@@ -259,6 +261,14 @@
# be an EagerTensor once again, not a graph Tensor.
self.assertEqual(float(step()), -1.0)
+ def testConstructSGDWithLR(self):
+ opt = gradient_descent.SGD(lr=1.0)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = gradient_descent.SGD(learning_rate=0.1, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = gradient_descent.SGD(learning_rate=0.1)
+ self.assertEqual(opt_3.lr, 0.1)
+
class MomentumOptimizerTest(test.TestCase):
@@ -346,7 +356,7 @@
var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
accum0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
accum1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
- loss = 5 * var0 * var0 + 3 * var1
+ loss = lambda: 5 * var0 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop
mom_op = gradient_descent.SGD(
learning_rate=2.0, momentum=0.9, nesterov=True)
opt_op = mom_op.minimize(loss, [var0, var1])
@@ -677,12 +687,20 @@
opt3._get_hyper("momentum"))
# self.assertEqual(
# self.evaluate(opt._get_hyper("decay")), opt3._get_hyper("decay"))
- self.assertTrue(opt3._nesterov)
+ self.assertTrue(opt3.nesterov)
def testNesterovWithoutMomentum(self):
with self.assertRaisesRegexp(ValueError, "must be between"):
gradient_descent.SGD(learning_rate=1.0, momentum=2.0)
+ def testConstructMomentumWithLR(self):
+ opt = gradient_descent.SGD(lr=1.0, momentum=0.9)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = gradient_descent.SGD(learning_rate=0.1, momentum=0.9, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = gradient_descent.SGD(learning_rate=0.1, momentum=0.9)
+ self.assertEqual(opt_3.lr, 0.1)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizer_v2/nadam.py b/tensorflow/python/keras/optimizer_v2/nadam.py
index 00b095e..afa74c8 100644
--- a/tensorflow/python/keras/optimizer_v2/nadam.py
+++ b/tensorflow/python/keras/optimizer_v2/nadam.py
@@ -74,6 +74,9 @@
**kwargs: keyword arguments. Allowed to be {`decay`}
"""
+ # Backwards compatiblity with keras NAdam optimizer.
+ if 'schedule_decay' in kwargs:
+ kwargs['decay'] = kwargs.pop('schedule_decay')
# pylint: disable=useless-super-delegation
super(Nadam, self).__init__(
learning_rate=learning_rate,
diff --git a/tensorflow/python/keras/optimizer_v2/nadam_test.py b/tensorflow/python/keras/optimizer_v2/nadam_test.py
index d991e31..73568e8 100644
--- a/tensorflow/python/keras/optimizer_v2/nadam_test.py
+++ b/tensorflow/python/keras/optimizer_v2/nadam_test.py
@@ -208,6 +208,18 @@
self.assertAllCloseAccordingToType(var0_np, var0.eval())
self.assertAllCloseAccordingToType(var1_np, var1.eval())
+ def testConstructNAdamWithLR(self):
+ opt = nadam.Nadam(lr=1.0)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = nadam.Nadam(learning_rate=0.1, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = nadam.Nadam(learning_rate=0.1)
+ self.assertEqual(opt_3.lr, 0.1)
+
+ def testConstructNAdamWithScheduleDecay(self):
+ opt = nadam.Nadam(schedule_decay=0.2)
+ self.assertEqual(opt.decay, 0.2)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
index b26b3ce..874d0f7 100644
--- a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
@@ -21,9 +21,12 @@
from __future__ import print_function
import abc
+import functools
import six
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribution_strategy_context as distribute_ctx
from tensorflow.python.distribute import reduce_util as ds_reduce_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
@@ -32,17 +35,40 @@
from tensorflow.python.keras import backend
from tensorflow.python.keras import initializers
from tensorflow.python.keras.engine import base_layer_utils
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
from tensorflow.python.ops import gradients
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables as tf_variables
from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.training import distribution_strategy_context
-from tensorflow.python.training import optimizer as optimizer_v1
+from tensorflow.python.training.checkpointable import base as checkpointable
from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _deduplicate_indexed_slices(values, indices):
+ """Sums `values` associated with any non-unique `indices`.
+
+ Args:
+ values: A `Tensor` with rank >= 1.
+ indices: A one-dimensional integer `Tensor`, indexing into the first
+ dimension of `values` (as in an IndexedSlices object).
+
+ Returns:
+ A tuple of (`summed_values`, `unique_indices`) where `unique_indices` is a
+ de-duplicated version of `indices` and `summed_values` contains the sum of
+ `values` slices associated with each unique index.
+ """
+ unique_indices, new_index_positions = array_ops.unique(indices)
+ summed_values = math_ops.unsorted_segment_sum(
+ values, new_index_positions,
+ array_ops.shape(unique_indices)[0])
+ return (summed_values, unique_indices)
@six.add_metaclass(abc.ABCMeta)
-class OptimizerV2(optimizer_v1.Optimizer):
+@tf_export("keras.optimizers.Optimizer")
+class OptimizerV2(checkpointable.CheckpointableBase):
"""Updated base class for optimizers.
This class defines the API to add Ops to train a model. You never use this
@@ -137,26 +163,30 @@
_create_vars.
"""
self._use_locking = True
- super(OptimizerV2, self).__init__(self._use_locking, name)
+ self._name = name
self._hyper = {}
# dict: {variable name : {slot name : variable}}
self._slots = {}
+ self._slot_names = []
self._weights = []
+ # For implementing Checkpointable. Stores information about how to restore
+ # slot variables which have not yet been created
+ # (checkpointable._CheckpointPosition objects).
+ # {slot_name :
+ # {_var_key(variable_to_train): [checkpoint_position, ... ], ... },
+ # ... }
+ self._deferred_slot_restorations = {}
+
decay = kwargs.pop("decay", 0.0)
if decay < 0.:
raise ValueError("decay cannot be less than 0: {}".format(decay))
self._initial_decay = decay
+ self.__dict__.update(kwargs)
self._prepared = False
- def minimize(self,
- loss,
- var_list,
- aggregation_method=None,
- colocate_gradients_with_ops=False,
- name=None,
- grad_loss=None):
+ def minimize(self, loss, var_list, grad_loss=None, name=None):
"""Add operations to minimize `loss` by updating `var_list`.
This method simply combines calls `compute_gradients()` and
@@ -165,15 +195,11 @@
of using this function.
Args:
- loss: A `Tensor` containing the value to minimize.
+ loss: A callable taking no arguments which returns the value to minimize.
var_list: list or tuple of `Variable` objects to update to minimize
`loss`.
- aggregation_method: Specifies the method used to combine gradient terms.
- Valid values are defined in the class `AggregationMethod`.
- colocate_gradients_with_ops: If True, try colocating gradients with the
- corresponding op.
- name: Optional name for the returned operation.
grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.
+ name: Optional name for the returned operation.
Returns:
An Operation that updates the variables in `var_list`. If `global_step`
@@ -185,29 +211,16 @@
@compatibility(eager)
When eager execution is enabled, `loss` should be a Python function that
takes no arguments and computes the value to be minimized. Minimization (and
- gradient computation) is done with respect to the elements of `var_list` if
- not None, else with respect to any trainable variables created during the
- execution of the `loss` function. `gate_gradients`, `aggregation_method`,
- `colocate_gradients_with_ops` and `grad_loss` are ignored when eager
- execution is enabled.
+ gradient computation) is done with respect to the elements of `var_list`.
+ `grad_loss` is ignored when eager execution is enabled.
@end_compatibility
"""
- grads_and_vars = self.compute_gradients(
- loss,
- var_list=var_list,
- aggregation_method=aggregation_method,
- colocate_gradients_with_ops=colocate_gradients_with_ops,
- grad_loss=grad_loss)
+ grads_and_vars = self._compute_gradients(
+ loss, var_list=var_list, grad_loss=grad_loss)
return self.apply_gradients(grads_and_vars, name=name)
- def compute_gradients(self,
- loss,
- var_list,
- aggregation_method=None,
- colocate_gradients_with_ops=False,
- grad_loss=None,
- stop_gradients=None):
+ def _compute_gradients(self, loss, var_list, grad_loss=None):
"""Compute gradients of `loss` for the variables in `var_list`.
This is the first part of `minimize()`. It returns a list
@@ -217,19 +230,11 @@
given variable.
Args:
- loss: A Tensor containing the value to minimize or a callable taking no
- arguments which returns the value to minimize. When eager execution is
- enabled it must be a callable.
- var_list: Optional list or tuple of `tf.Variable` to update to minimize
+ loss: A callable taking no arguments which returns the value to minimize.
+ var_list: List or tuple of `tf.Variable` to update to minimize
`loss`. Defaults to the list of variables collected in the graph under
the key `GraphKeys.TRAINABLE_VARIABLES`.
- aggregation_method: Specifies the method used to combine gradient terms.
- Valid values are defined in the class `AggregationMethod`.
- colocate_gradients_with_ops: If True, try colocating gradients with the
- corresponding op.
grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.
- stop_gradients: Optional. A Tensor or list of tensors not to differentiate
- through.
Returns:
A list of (gradient, variable) pairs. Variable is always present, but
@@ -238,36 +243,22 @@
Raises:
TypeError: If `var_list` contains anything else than `Variable` objects.
ValueError: If some arguments are invalid, or var_list is None.
- RuntimeError: If called with eager execution enabled and `loss` is
- not callable.
-
- @compatibility(eager)
- When eager execution is enabled, `aggregation_method`, and
- `colocate_gradients_with_ops` are ignored.
- @end_compatibility
"""
var_list = nest.flatten(var_list)
# TODO(josh11b): Test that we handle weight decay in a reasonable way.
- if callable(loss):
- with backprop.GradientTape() as tape:
- tape.watch(var_list)
- loss_value = loss()
- grads = tape.gradient(loss_value, var_list, grad_loss)
- else:
- if context.executing_eagerly():
- raise RuntimeError("`loss` passed to Optimizer.compute_gradients "
- "should be a function when eager execution is "
- "enabled.")
- self._assert_valid_dtypes([loss])
- if grad_loss is not None:
- self._assert_valid_dtypes([grad_loss])
- grads = gradients.gradients(
- loss,
- var_list,
- grad_ys=grad_loss,
- aggregation_method=aggregation_method,
- colocate_gradients_with_ops=colocate_gradients_with_ops,
- stop_gradients=stop_gradients)
+ with backprop.GradientTape() as tape:
+ tape.watch(var_list)
+ loss_value = loss()
+ loss_value = self._scale_loss(loss_value)
+ grads = tape.gradient(loss_value, var_list, grad_loss)
+
+ if hasattr(self, "clipnorm"):
+ grads = [clip_ops.clip_by_norm(g, self.clipnorm) for g in grads]
+ if hasattr(self, "clipvalue"):
+ grads = [
+ clip_ops.clip_by_value(g, -self.clipvalue, self.clipvalue)
+ for g in grads
+ ]
grads_and_vars = list(zip(grads, var_list))
self._assert_valid_dtypes([
@@ -277,6 +268,46 @@
return grads_and_vars
+ @staticmethod
+ def _scale_loss(loss_value):
+ if distribute_lib.get_loss_reduction() == ds_reduce_util.ReduceOp.MEAN:
+ num_replicas = \
+ distribute_ctx.get_distribution_strategy().num_replicas_in_sync
+ if num_replicas > 1:
+ loss_value *= (1. / num_replicas)
+ return loss_value
+
+ def get_gradients(self, loss, params):
+ """Returns gradients of `loss` with respect to `params`.
+
+ Arguments:
+ loss: Loss tensor.
+ params: List of variables.
+
+ Returns:
+ List of gradient tensors.
+
+ Raises:
+ ValueError: In case any gradient cannot be computed (e.g. if gradient
+ function not implemented).
+ """
+ loss = self._scale_loss(loss)
+ grads = gradients.gradients(loss, params)
+ if None in grads:
+ raise ValueError("An operation has `None` for gradient. "
+ "Please make sure that all of your ops have a "
+ "gradient defined (i.e. are differentiable). "
+ "Common ops without gradient: "
+ "K.argmax, K.round, K.eval.")
+ if hasattr(self, "clipnorm"):
+ grads = [clip_ops.clip_by_norm(g, self.clipnorm) for g in grads]
+ if hasattr(self, "clipvalue"):
+ grads = [
+ clip_ops.clip_by_value(g, -self.clipvalue, self.clipvalue)
+ for g in grads
+ ]
+ return grads
+
def apply_gradients(self, grads_and_vars, name=None):
"""Apply gradients to variables.
@@ -299,7 +330,7 @@
"""
grads_and_vars = _filter_grads(grads_and_vars)
var_list = [v for (_, v) in grads_and_vars]
- if distribution_strategy_context.has_distribution_strategy():
+ if distribute_ctx.has_distribution_strategy():
reduced_grads = merge_grads(grads_and_vars)
grads_and_vars = zip(reduced_grads, var_list)
@@ -339,7 +370,13 @@
return apply_updates
def get_updates(self, loss, params):
- return [self.minimize(loss, params)]
+ grads = self.get_gradients(loss, params)
+ grads_and_vars = list(zip(grads, params))
+ self._assert_valid_dtypes([
+ v for g, v in grads_and_vars
+ if g is not None and v.dtype != dtypes.resource
+ ])
+ return [self.apply_gradients(grads_and_vars)]
def _set_hyper(self, name, value):
"""set hyper `name` to value. value can be callable, tensor, numeric."""
@@ -387,18 +424,36 @@
else:
super(OptimizerV2, self).__setattr__(name, value)
+ def get_slot_names(self):
+ """A list of names for this optimizer's slots."""
+ return self._slot_names
+
def add_slot(self, var, slot_name, initializer="zeros"):
+ """Add a new slot variable for `var`."""
+ if slot_name not in self._slot_names:
+ self._slot_names.append(slot_name)
var_key = _var_key(var)
slot_dict = self._slots.setdefault(var_key, {})
- if slot_name not in slot_dict:
- slot_key = _get_slot_key_from_var(var, slot_name)
- weight = self.add_weight(
- name=slot_key,
- shape=var.shape,
+ weight = slot_dict.get(slot_name, None)
+ if weight is None:
+ if isinstance(initializer, six.string_types) or callable(initializer):
+ initializer = initializers.get(initializer)
+ initial_value = functools.partial(
+ initializer, shape=var.shape, dtype=var.dtype)
+ else:
+ initial_value = initializer
+ weight = tf_variables.Variable(
+ name="%s/%s" % (var._shared_name, slot_name), # pylint: disable=protected-access
dtype=var.dtype,
- initializer=initializer)
+ trainable=False,
+ initial_value=initial_value)
+ backend.track_variable(weight)
slot_dict[slot_name] = weight
+ self._restore_slot_variable(
+ slot_name=slot_name, variable=var,
+ slot_variable=weight)
self._weights.append(weight)
+ return weight
def get_slot(self, var, slot_name):
var_key = _var_key(var)
@@ -455,7 +510,12 @@
Returns:
Python dictionary.
"""
- return {"name": self._name}
+ config = {"name": self._name}
+ if hasattr(self, "clipnorm"):
+ config["clipnorm"] = self.clipnorm
+ if hasattr(self, "clipvalue"):
+ config["clipvalue"] = self.clipvalue
+ return config
@classmethod
def from_config(cls, config, custom_objects=None):
@@ -563,12 +623,178 @@
return variable
+ def _assert_valid_dtypes(self, tensors):
+ """Asserts tensors are all valid types (see `_valid_dtypes`).
+
+ Args:
+ tensors: Tensors to check.
+
+ Raises:
+ ValueError: If any tensor is not a valid type.
+ """
+ valid_dtypes = self._valid_dtypes()
+ for t in tensors:
+ dtype = t.dtype.base_dtype
+ if dtype not in valid_dtypes:
+ raise ValueError("Invalid type %r for %s, expected: %s." %
+ (dtype, t.name, [v for v in valid_dtypes]))
+
+ def _valid_dtypes(self):
+ """Valid types for loss, variables and gradients.
+
+ Subclasses should override to allow other float types.
+
+ Returns:
+ Valid types for loss, variables and gradients.
+ """
+ return set(
+ [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64])
+
+ def _call_if_callable(self, param):
+ """Call the function if param is callable."""
+ return param() if callable(param) else param
+
+ def _resource_apply_dense(self, grad, handle):
+ """Add ops to apply dense gradients to the variable `handle`.
+
+ Args:
+ grad: a `Tensor` representing the gradient.
+ handle: a `Tensor` of dtype `resource` which points to the variable to be
+ updated.
+
+ Returns:
+ An `Operation` which updates the value of the variable.
+ """
+ raise NotImplementedError()
+
+ def _resource_apply_sparse_duplicate_indices(self, grad, handle, indices):
+ """Add ops to apply sparse gradients to `handle`, with repeated indices.
+
+ Optimizers which override this method must deal with repeated indices. See
+ the docstring of `_apply_sparse_duplicate_indices` for details. By default
+ the correct behavior, to sum non-unique indices and their associated
+ gradients, is enforced by first pre-processing `grad` and `indices` and
+ passing them on to `_resource_apply_sparse`. Optimizers which deal correctly
+ with duplicate indices may instead override this method to avoid the
+ overhead of summing.
+
+ Args:
+ grad: a `Tensor` representing the gradient for the affected indices.
+ handle: a `Tensor` of dtype `resource` which points to the variable to be
+ updated.
+ indices: a `Tensor` of integral type representing the indices for which
+ the gradient is nonzero. Indices may be repeated.
+
+ Returns:
+ An `Operation` which updates the value of the variable.
+ """
+ summed_grad, unique_indices = _deduplicate_indexed_slices(
+ values=grad, indices=indices)
+ return self._resource_apply_sparse(summed_grad, handle, unique_indices)
+
+ def _resource_apply_sparse(self, grad, handle, indices):
+ """Add ops to apply sparse gradients to the variable `handle`.
+
+ Similar to `_apply_sparse`, the `indices` argument to this method has been
+ de-duplicated. Optimizers which deal correctly with non-unique indices may
+ instead override `_resource_apply_sparse_duplicate_indices` to avoid this
+ overhead.
+
+ Args:
+ grad: a `Tensor` representing the gradient for the affected indices.
+ handle: a `Tensor` of dtype `resource` which points to the variable to be
+ updated.
+ indices: a `Tensor` of integral type representing the indices for which
+ the gradient is nonzero. Indices are unique.
+
+ Returns:
+ An `Operation` which updates the value of the variable.
+ """
+ raise NotImplementedError()
+
+ # ---------------
+ # For implementing the checkpointable interface
+ # ---------------
+
+ def _restore_slot_variable(self, slot_name, variable, slot_variable):
+ """Restore a newly created slot variable's value."""
+ variable_key = _var_key(variable)
+ deferred_restorations = self._deferred_slot_restorations.get(
+ slot_name, {}).pop(variable_key, [])
+ # Iterate over restores, highest restore UID first to minimize the number
+ # of assignments.
+ deferred_restorations.sort(key=lambda position: position.restore_uid,
+ reverse=True)
+ for checkpoint_position in deferred_restorations:
+ checkpoint_position.restore(slot_variable)
+
+ def _create_or_restore_slot_variable(
+ self, slot_variable_position, slot_name, variable):
+ """Restore a slot variable's value, possibly creating it.
+
+ Called when a variable which has an associated slot variable is created or
+ restored. When executing eagerly, we create the slot variable with a
+ restoring initializer.
+
+ No new variables are created when graph building. Instead,
+ _restore_slot_variable catches these after normal creation and adds restore
+ ops to the graph. This method is nonetheless important when graph building
+ for the case when a slot variable has already been created but `variable`
+ has just been added to a dependency graph (causing us to realize that the
+ slot variable needs to be restored).
+
+ Args:
+ slot_variable_position: A `checkpointable._CheckpointPosition` object
+ indicating the slot variable `Checkpointable` object to be restored.
+ slot_name: The name of this `Optimizer`'s slot to restore into.
+ variable: The variable object this slot is being created for.
+ """
+ variable_key = _var_key(variable)
+ slot_dict = self._slots.get(variable_key, {})
+ slot_variable = slot_dict.get(slot_name, None)
+ if (slot_variable is None and context.executing_eagerly() and
+ slot_variable_position.is_simple_variable()
+ # Defer slot variable creation if there is an active variable creator
+ # scope. Generally we'd like to eagerly create/restore slot variables
+ # when possible, but this may mean that scopes intended to catch
+ # `variable` also catch its eagerly created slot variable
+ # unintentionally (specifically make_template would add a dependency on
+ # a slot variable if not for this case). Deferring is mostly harmless
+ # (aside from double initialization), and makes variable creator scopes
+ # behave the same way they do when graph building.
+ and not ops.get_default_graph()._variable_creator_stack): # pylint: disable=protected-access
+ initializer = checkpointable.CheckpointInitialValue(
+ checkpoint_position=slot_variable_position)
+ slot_variable = self.add_slot(
+ var=variable,
+ initializer=initializer,
+ slot_name=slot_name)
+ # Slot variables are not owned by any one object (because we don't want to
+ # save the slot variable if the optimizer is saved without the non-slot
+ # variable, or if the non-slot variable is saved without the optimizer;
+ # it's a dependency hypergraph with edges of the form (optimizer, non-slot
+ # variable, variable)). So we don't _track_ slot variables anywhere, and
+ # instead special-case this dependency and otherwise pretend it's a normal
+ # graph.
+ if slot_variable is not None:
+ # If we've either made this slot variable, or if we've pulled out an
+ # existing slot variable, we should restore it.
+ slot_variable_position.restore(slot_variable)
+ else:
+ # We didn't make the slot variable. Defer restoring until it gets created
+ # normally. We keep a list rather than the one with the highest restore
+ # UID in case slot variables have their own dependencies, in which case
+ # those could differ between restores.
+ self._deferred_slot_restorations.setdefault(
+ slot_name, {}).setdefault(variable_key, []).append(
+ slot_variable_position)
+
def _filter_grads(grads_and_vars):
"""Filter out iterable with grad equal to None."""
grads_and_vars = tuple(grads_and_vars)
if not grads_and_vars:
- raise ValueError("No variables provided.")
+ return grads_and_vars
filtered = []
vars_with_empty_grads = []
for grad, var in grads_and_vars:
@@ -598,7 +824,7 @@
incre_op = local_step.assign_add(1).op
return incre_op
- return distribution_strategy_context.get_replica_context().merge_call(
+ return distribute_ctx.get_replica_context().merge_call(
merge_update_step_fn, args=(update_ops, local_step))
@@ -606,11 +832,11 @@
"""Merge gradients from different replicas."""
def merge_grad_fn(strategy, grads_and_vars):
- reduced_grads = strategy.batch_reduce(
- ds_reduce_util.ReduceOp.MEAN, grads_and_vars)
+ reduced_grads = strategy.batch_reduce(ds_reduce_util.ReduceOp.SUM,
+ grads_and_vars)
return reduced_grads
- return distribution_strategy_context.get_replica_context().merge_call(
+ return distribute_ctx.get_replica_context().merge_call(
merge_grad_fn, args=(grads_and_vars,))
@@ -629,7 +855,7 @@
"""
# pylint: disable=protected-access
- if distribution_strategy_context.has_distribution_strategy() and hasattr(
+ if distribute_ctx.has_distribution_strategy() and hasattr(
var, "_primary_var"):
var = var._primary_var
if hasattr(var, "op"):
diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2_test.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2_test.py
index 158577f..8b2865e 100644
--- a/tensorflow/python/keras/optimizer_v2/optimizer_v2_test.py
+++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2_test.py
@@ -46,7 +46,6 @@
from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import clip_ops
-from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables
@@ -64,8 +63,6 @@
var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
loss = lambda: 5 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop
- if not context.executing_eagerly():
- loss = loss()
sgd = gradient_descent.SGD(3.0)
self.evaluate(variables.global_variables_initializer())
@@ -117,41 +114,12 @@
self.assertAllClose([-1.5, -0.5], self.evaluate(var1))
@test_util.run_in_graph_and_eager_modes
- def testAggregationMethod(self):
- for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
- with self.cached_session():
- var0 = variables.Variable([1.0, 2.0], dtype=dtype)
- var1 = variables.Variable([3.0, 4.0], dtype=dtype)
- loss = lambda: 5 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop
- if not context.executing_eagerly():
- loss = loss()
- sgd = gradient_descent.SGD(3.0)
-
- self.evaluate(variables.global_variables_initializer())
- # Fetch params to validate initial values
- self.assertAllClose([1.0, 2.0], self.evaluate(var0))
- self.assertAllClose([3.0, 4.0], self.evaluate(var1))
- # Run 1 step of sgd through optimizer
- opt_op = sgd.minimize(
- loss,
- var_list=[var0, var1],
- aggregation_method=gradients_impl.AggregationMethod
- .EXPERIMENTAL_ACCUMULATE_N)
- self.evaluate(variables.global_variables_initializer())
- self.evaluate(opt_op)
- # Validate updated params
- self.assertAllClose([-14., -13.], self.evaluate(var0))
- self.assertAllClose([-6., -5.], self.evaluate(var1))
-
- @test_util.run_in_graph_and_eager_modes
def testPrecomputedGradient(self):
for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
with self.cached_session():
var0 = variables.Variable([1.0, 2.0], dtype=dtype)
var1 = variables.Variable([3.0, 4.0], dtype=dtype)
loss = lambda: 5 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop
- if not context.executing_eagerly():
- loss = loss()
grad_loss = constant_op.constant([42, -42], dtype=dtype)
sgd = gradient_descent.SGD(3.0)
@@ -176,8 +144,6 @@
var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
loss = lambda: 5 * var0 # pylint: disable=cell-var-from-loop
- if not context.executing_eagerly():
- loss = loss()
sgd_op = gradient_descent.SGD(3.0)
with self.assertRaisesRegexp(ValueError, 'No gradients'):
# var1 has no gradient
@@ -190,8 +156,6 @@
var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
loss = lambda: constant_op.constant(5.0)
- if not context.executing_eagerly():
- loss = loss()
sgd_op = gradient_descent.SGD(3.0)
with self.assertRaisesRegexp(ValueError,
@@ -216,11 +180,9 @@
var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
loss = lambda: 5 * var0 + 3 * var1 # pylint: disable=cell-var-from-loop
- if not context.executing_eagerly():
- loss = loss()
sgd = gradient_descent.SGD(3.0)
- grads_and_vars = sgd.compute_gradients(loss, [var0, var1])
+ grads_and_vars = sgd._compute_gradients(loss, [var0, var1])
# Convert gradients to tf.Variables
converted_grads = [
resource_variable_ops.ResourceVariable(
@@ -259,7 +221,7 @@
return x * x
sgd = gradient_descent.SGD(3.0)
- grads_and_vars = sgd.compute_gradients(f, [x])
+ grads_and_vars = sgd._compute_gradients(f, [x])
self.assertEqual(1, len(grads_and_vars))
grad, x_as_var = grads_and_vars[0]
self.assertIs(x, x_as_var)
@@ -278,8 +240,6 @@
var1 = variables.Variable([3.0, 4.0],
constraint=constraint_0)
loss = lambda: 5 * var0 + 3 * var1
- if not context.executing_eagerly(): # pylint: disable=cell-var-from-loop
- loss = loss()
sgd = gradient_descent.SGD(3.0)
self.evaluate(variables.global_variables_initializer())
@@ -339,6 +299,28 @@
opt3._get_hyper('learning_rate'))
@test_util.run_in_graph_and_eager_modes
+ def testGradClipValue(self):
+ with self.cached_session():
+ var = resource_variable_ops.ResourceVariable([1.0, 2.0])
+ loss = lambda: 3 * var
+ opt = gradient_descent.SGD(learning_rate=1.0, clipvalue=1.0)
+ opt_op = opt.minimize(loss, [var])
+ self.evaluate(variables.global_variables_initializer())
+ self.evaluate(opt_op)
+ self.assertAllClose([0., 1.], self.evaluate(var))
+
+ @test_util.run_in_graph_and_eager_modes
+ def testGradClipNorm(self):
+ with self.cached_session():
+ var = resource_variable_ops.ResourceVariable([1.0])
+ loss = lambda: 3 * var
+ opt = gradient_descent.SGD(learning_rate=1.0, clipnorm=1.0)
+ opt_op = opt.minimize(loss, [var])
+ self.evaluate(variables.global_variables_initializer())
+ self.evaluate(opt_op)
+ self.assertAllClose([0.], self.evaluate(var))
+
+ @test_util.run_in_graph_and_eager_modes
def testWeights(self):
with self.cached_session():
opt1 = adam.Adam(learning_rate=1.0)
diff --git a/tensorflow/python/keras/optimizer_v2/rmsprop.py b/tensorflow/python/keras/optimizer_v2/rmsprop.py
index 6a5b334..634111b 100644
--- a/tensorflow/python/keras/optimizer_v2/rmsprop.py
+++ b/tensorflow/python/keras/optimizer_v2/rmsprop.py
@@ -20,8 +20,10 @@
from tensorflow.python.framework import ops
from tensorflow.python.keras.optimizer_v2 import optimizer_v2
from tensorflow.python.training import training_ops
+from tensorflow.python.util.tf_export import tf_export
+@tf_export("keras.optimizers.RMSprop")
class RMSprop(optimizer_v2.OptimizerV2):
r"""Optimizer that implements the RMSprop algorithm.
@@ -91,7 +93,7 @@
**kwargs: keyword arguments. Allowed to be {`decay`}
"""
super(RMSprop, self).__init__(name, **kwargs)
- self._set_hyper("learning_rate", learning_rate)
+ self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
self._set_hyper("decay", self._initial_decay)
self._set_hyper("rho", rho)
@@ -103,13 +105,13 @@
self._set_hyper("momentum", momentum)
self._set_hyper("epsilon", epsilon)
- self._centered = centered
+ self.centered = centered
def _create_slots(self, var_list):
for var in var_list:
self.add_slot(var, "rms")
self.add_slot(var, "momentum")
- if self._centered:
+ if self.centered:
self.add_slot(var, "mg")
def _resource_apply_dense(self, grad, var):
@@ -120,7 +122,7 @@
rho = self._get_hyper("rho", var_dtype)
momentum = self._get_hyper("momentum", var_dtype)
epsilon = self._get_hyper("epsilon", var_dtype)
- if self._centered:
+ if self.centered:
mg = self.get_slot(var, "mg")
return training_ops.resource_apply_centered_rms_prop(
var.handle,
@@ -153,7 +155,7 @@
rho = self._get_hyper("rho", var_dtype)
momentum = self._get_hyper("momentum", var_dtype)
epsilon = self._get_hyper("epsilon", var_dtype)
- if self._centered:
+ if self.centered:
mg = self.get_slot(var, "mg")
return training_ops.resource_sparse_apply_centered_rms_prop(
var.handle,
@@ -188,7 +190,7 @@
"rho": self._serialize_hyperparameter("rho"),
"momentum": self._serialize_hyperparameter("momentum"),
"epsilon": self._serialize_hyperparameter("epsilon"),
- "centered": self._centered,
+ "centered": self.centered,
})
return config
diff --git a/tensorflow/python/keras/optimizer_v2/rmsprop_test.py b/tensorflow/python/keras/optimizer_v2/rmsprop_test.py
index a8658a8..4d61cfb 100644
--- a/tensorflow/python/keras/optimizer_v2/rmsprop_test.py
+++ b/tensorflow/python/keras/optimizer_v2/rmsprop_test.py
@@ -233,8 +233,11 @@
with self.cached_session():
var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
- pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
- loss = pred * pred
+
+ def loss():
+ pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x) # pylint: disable=cell-var-from-loop
+ return pred * pred
+
sgd_op = rmsprop.RMSprop(
learning_rate=1.0,
rho=0.0,
@@ -258,8 +261,12 @@
with self.cached_session():
var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype)
x = constant_op.constant([[4.0], [5.0]], dtype=dtype)
- pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
- loss = pred * pred
+
+ def loss():
+ pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x) # pylint: disable=cell-var-from-loop
+ return pred * pred
+
+ # loss = lambda: pred * pred # pylint: disable=cell-var-from-loop
sgd_op = rmsprop.RMSprop(
learning_rate=1.0,
rho=0.0,
@@ -405,6 +412,14 @@
(0.01 * 2.0 / math.sqrt(0.00001 * 0.9 + 1e-5 + 1.0))
]), self.evaluate(var1))
+ def testConstructRMSpropWithLR(self):
+ opt = rmsprop.RMSprop(lr=1.0)
+ self.assertEqual(opt.lr, 1.0)
+ opt_2 = rmsprop.RMSprop(learning_rate=0.1, lr=1.0)
+ self.assertEqual(opt_2.lr, 1.0)
+ opt_3 = rmsprop.RMSprop(learning_rate=0.1)
+ self.assertEqual(opt_3.lr, 0.1)
+
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/keras/optimizers.py b/tensorflow/python/keras/optimizers.py
index 10466eb..decfcf9 100644
--- a/tensorflow/python/keras/optimizers.py
+++ b/tensorflow/python/keras/optimizers.py
@@ -23,6 +23,7 @@
from six.moves import zip # pylint: disable=redefined-builtin
from tensorflow.python import tf2
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.framework import ops
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.optimizer_v2 import adadelta as adadelta_v2
@@ -38,14 +39,12 @@
from tensorflow.python.ops import clip_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
-from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.training import optimizer as tf_optimizer_module
from tensorflow.python.training import training_util
from tensorflow.python.training.checkpointable import base as checkpointable
from tensorflow.python.util.tf_export import tf_export
-@tf_export('keras.optimizers.Optimizer')
class Optimizer(object):
"""Abstract optimizer base class.
@@ -159,7 +158,6 @@
return cls(**config)
-@tf_export('keras.optimizers.SGD')
class SGD(Optimizer):
"""Stochastic gradient descent optimizer.
@@ -224,7 +222,6 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('keras.optimizers.RMSprop')
class RMSprop(Optimizer):
"""RMSProp optimizer.
@@ -291,7 +288,6 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('keras.optimizers.Adagrad')
class Adagrad(Optimizer):
"""Adagrad optimizer.
@@ -358,7 +354,6 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('keras.optimizers.Adadelta')
class Adadelta(Optimizer):
"""Adadelta optimizer.
@@ -442,7 +437,6 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('keras.optimizers.Adam')
class Adam(Optimizer):
"""Adam optimizer.
@@ -539,7 +533,6 @@
return dict(list(base_config.items()) + list(config.items()))
-@tf_export('keras.optimizers.Adamax')
class Adamax(Optimizer):
"""Adamax optimizer from Adam paper's Section 7.
@@ -806,27 +799,18 @@
Returns:
A Keras Optimizer instance.
"""
- if tf2.enabled():
- all_classes = {
- 'adadelta': adadelta_v2.Adadelta,
- 'adagrad': adagrad_v2.Adagrad,
- 'adam': adam_v2.Adam,
- 'adamax': adamax_v2.Adamax,
- 'nadam': nadam_v2.Nadam,
- 'rmsprop': rmsprop_v2.RMSprop,
- 'sgd': gradient_descent_v2.SGD
- }
- else:
- all_classes = {
- 'adadelta': Adadelta,
- 'adagrad': Adagrad,
- 'adam': Adam,
- 'adamax': Adamax,
- 'nadam': Nadam,
- 'rmsprop': RMSprop,
- 'sgd': SGD,
- 'tfoptimizer': TFOptimizer
- }
+ all_classes = {
+ 'adadelta': adadelta_v2.Adadelta,
+ 'adagrad': adagrad_v2.Adagrad,
+ 'adam': adam_v2.Adam,
+ 'adamax': adamax_v2.Adamax,
+ 'nadam': nadam_v2.Nadam,
+ 'rmsprop': rmsprop_v2.RMSprop,
+ 'sgd': gradient_descent_v2.SGD
+ }
+ if not tf2.enabled():
+ all_classes['nadam'] = Nadam
+
# Make deserialization case-insensitive for built-in optimizers.
if config['class_name'].lower() in all_classes:
config['class_name'] = config['class_name'].lower()
diff --git a/tensorflow/python/keras/optimizers_test.py b/tensorflow/python/keras/optimizers_test.py
index d3cacb7..33d65e6 100644
--- a/tensorflow/python/keras/optimizers_test.py
+++ b/tensorflow/python/keras/optimizers_test.py
@@ -65,6 +65,15 @@
optim = keras.optimizers.deserialize(config)
new_config = keras.optimizers.serialize(optim)
new_config['class_name'] = new_config['class_name'].lower()
+ new_config['config'].pop('name', None)
+ if 'amsgrad' not in config['config']:
+ new_config['config'].pop('amsgrad', None)
+ if 'decay' in new_config['config'] and 'schedule_decay' in config['config']:
+ new_config['config']['schedule_decay'] = new_config['config'].pop('decay')
+ if 'momentum' not in config['config']:
+ new_config['config'].pop('momentum', None)
+ if 'centered' not in config['config']:
+ new_config['config'].pop('centered', None)
assert config == new_config
# Test constraints.
@@ -91,26 +100,22 @@
class KerasOptimizersTest(test.TestCase):
- @test_util.run_deprecated_v1
def test_sgd(self):
with self.cached_session():
_test_optimizer(keras.optimizers.SGD(lr=0.01,
momentum=0.9,
nesterov=True))
- @test_util.run_deprecated_v1
def test_rmsprop(self):
with self.cached_session():
_test_optimizer(keras.optimizers.RMSprop())
_test_optimizer(keras.optimizers.RMSprop(decay=1e-3))
- @test_util.run_deprecated_v1
def test_adagrad(self):
with self.cached_session():
_test_optimizer(keras.optimizers.Adagrad())
_test_optimizer(keras.optimizers.Adagrad(decay=1e-3))
- @test_util.run_deprecated_v1
def test_adadelta(self):
with self.cached_session():
_test_optimizer(keras.optimizers.Adadelta(), target=0.6)
@@ -119,32 +124,27 @@
# the accuracy.
_test_optimizer(keras.optimizers.Adadelta(decay=1e-3), target=0.4)
- @test_util.run_deprecated_v1
def test_adam(self):
with self.cached_session():
_test_optimizer(keras.optimizers.Adam())
_test_optimizer(keras.optimizers.Adam(decay=1e-3))
_test_optimizer(keras.optimizers.Adam(amsgrad=True))
- @test_util.run_deprecated_v1
def test_adamax(self):
with self.cached_session():
_test_optimizer(keras.optimizers.Adamax())
_test_optimizer(keras.optimizers.Adamax(decay=1e-3))
- @test_util.run_deprecated_v1
def test_nadam(self):
with self.cached_session():
_test_optimizer(keras.optimizers.Nadam())
- @test_util.run_deprecated_v1
def test_clipnorm(self):
with self.cached_session():
_test_optimizer(keras.optimizers.SGD(lr=0.01,
momentum=0.9,
clipnorm=0.5))
- @test_util.run_deprecated_v1
def test_clipvalue(self):
with self.cached_session():
_test_optimizer(keras.optimizers.SGD(lr=0.01,
diff --git a/tensorflow/python/keras/testing_utils.py b/tensorflow/python/keras/testing_utils.py
index 59ed776..fd062b0 100644
--- a/tensorflow/python/keras/testing_utils.py
+++ b/tensorflow/python/keras/testing_utils.py
@@ -77,9 +77,13 @@
Returns:
The output data (Numpy array) returned by the layer, for additional
checks to be done by the calling code.
+
+ Raises:
+ ValueError: if `input_shape is None`.
"""
if input_data is None:
- assert input_shape
+ if input_shape is None:
+ raise ValueError('input_shape is None')
if not input_dtype:
input_dtype = 'float32'
input_data_shape = list(input_shape)
@@ -157,7 +161,11 @@
# train(). This was causing some error for layer with Defun as it body.
# See b/120160788 for more details. This should be mitigated after 2.0.
model = keras.models.Model(x, layer(x))
- model.compile(RMSPropOptimizer(0.01), 'mse', weighted_metrics=['acc'])
+ if _thread_local_data.run_eagerly is not None:
+ model.compile(RMSPropOptimizer(0.01), 'mse', weighted_metrics=['acc'],
+ run_eagerly=should_run_eagerly())
+ else:
+ model.compile(RMSPropOptimizer(0.01), 'mse', weighted_metrics=['acc'])
model.train_on_batch(input_data, actual_output)
# test as first layer in Sequential API
diff --git a/tensorflow/python/keras/utils/data_utils.py b/tensorflow/python/keras/utils/data_utils.py
index 01a9d61..d133e3f 100644
--- a/tensorflow/python/keras/utils/data_utils.py
+++ b/tensorflow/python/keras/utils/data_utils.py
@@ -596,9 +596,9 @@
Function, a Function to initialize the pool
"""
def pool_fn(seqs):
- return multiprocessing.Pool(workers,
- initializer=init_pool_generator,
- initargs=(seqs, self.random_seed))
+ return multiprocessing.Pool(
+ workers, initializer=init_pool_generator, initargs=(seqs, None))
+
return pool_fn
def _wait_queue(self):
diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD
index f990b0e..ddb2dda 100644
--- a/tensorflow/python/kernel_tests/BUILD
+++ b/tensorflow/python/kernel_tests/BUILD
@@ -272,7 +272,7 @@
cuda_py_test(
name = "ctc_loss_op_test",
- size = "small",
+ size = "medium",
srcs = ["ctc_loss_op_test.py"],
additional_deps = [
"//third_party/py/numpy",
@@ -1190,8 +1190,15 @@
srcs = ["unicode_decode_op_test.py"],
additional_deps = [
"@absl_py//absl/testing:parameterized",
+ "//tensorflow/python/ops/ragged:ragged_factory_ops",
+ "//tensorflow/python:array_ops",
"//tensorflow/python:client_testlib",
+ "//tensorflow/python:errors",
"//tensorflow/python:framework_for_generated_wrappers",
+ "//tensorflow/python/ops/ragged:ragged",
+ "//tensorflow/python/ops/ragged:ragged_test_util",
+ "//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:sparse_tensor",
"//tensorflow/python:string_ops",
],
)
@@ -1889,6 +1896,22 @@
)
cuda_py_test(
+ name = "lu_op_test",
+ size = "small",
+ srcs = ["lu_op_test.py"],
+ additional_deps = [
+ "//third_party/py/numpy",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:client_testlib",
+ "//tensorflow/python:framework_for_generated_wrappers",
+ "//tensorflow/python:linalg_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:platform",
+ "//tensorflow/python/ops/linalg",
+ ],
+)
+
+cuda_py_test(
name = "manip_ops_test",
size = "small",
srcs = ["manip_ops_test.py"],
diff --git a/tensorflow/python/kernel_tests/atrous_convolution_test.py b/tensorflow/python/kernel_tests/atrous_convolution_test.py
index 6b16fca..2fb8a37 100644
--- a/tensorflow/python/kernel_tests/atrous_convolution_test.py
+++ b/tensorflow/python/kernel_tests/atrous_convolution_test.py
@@ -110,6 +110,7 @@
add_check(check, y1, y2)
+ @test_util.run_v1_only("b/120545219")
def test_unknown_spatial_dims_for_channel_last_format(self):
x = array_ops.placeholder(dtypes.float32, [1, None, None, 10])
w = array_ops.zeros([3, 3, 10, 20])
@@ -117,6 +118,7 @@
x, w, "VALID", dilation_rate=[2, 2], data_format="NHWC")
self.assertEqual(y.shape.as_list(), [1, None, None, 20])
+ @test_util.run_v1_only("b/120545219")
def test_unknown_spatial_dims_for_channel_first_format(self):
x = array_ops.placeholder(dtypes.float32, [1, 10, None, None])
w = array_ops.zeros([3, 3, 10, 20])
@@ -262,6 +264,7 @@
err_tolerance = 1e-3
self.assertLess(err, err_tolerance)
+ @test_util.run_v1_only("b/120545219")
def testGradient(self):
with self.cached_session():
for padding in ["SAME", "VALID"]:
diff --git a/tensorflow/python/kernel_tests/base64_ops_test.py b/tensorflow/python/kernel_tests/base64_ops_test.py
index bb903d8..381f190 100644
--- a/tensorflow/python/kernel_tests/base64_ops_test.py
+++ b/tensorflow/python/kernel_tests/base64_ops_test.py
@@ -31,6 +31,7 @@
from tensorflow.python.platform import test
+@test_util.run_v1_only("b/120545219")
class Base64OpsTest(test_util.TensorFlowTestCase):
def setUp(self):
diff --git a/tensorflow/python/kernel_tests/batch_matmul_op_test.py b/tensorflow/python/kernel_tests/batch_matmul_op_test.py
index a0ad815..c32a6c7 100644
--- a/tensorflow/python/kernel_tests/batch_matmul_op_test.py
+++ b/tensorflow/python/kernel_tests/batch_matmul_op_test.py
@@ -20,9 +20,10 @@
import numpy as np
-from tensorflow.python.framework import constant_op
+from tensorflow.python import tf2
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
@@ -105,36 +106,37 @@
def _testNonEmpty(self, dtype, adjoint_a, adjoint_b, use_static_shape):
- def compareNonEmpty(self, a_shape, b_shape):
+ def CompareNonEmpty(self, a_shape, b_shape):
self._compare(
self._rand(a_shape, dtype),
self._rand(b_shape, dtype), adjoint_a, adjoint_b, use_static_shape)
- compareNonEmpty(self, [1, 2, 3], [1, 3, 5])
- compareNonEmpty(self, [1, 2, 3], [1, 3, 1])
- compareNonEmpty(self, [1, 1, 3], [1, 3, 5])
- compareNonEmpty(self, [1, 2, 3], [1, 3, 5])
- compareNonEmpty(self, [7, 1, 3], [7, 3, 5])
- compareNonEmpty(self, [7, 2, 3], [7, 3, 1])
- compareNonEmpty(self, [7, 2, 3], [7, 3, 5])
- compareNonEmpty(self, [10, 64, 75], [10, 75, 30])
- compareNonEmpty(self, [5, 7, 2, 3], [5, 7, 3, 5])
+ CompareNonEmpty(self, [1, 2, 3], [1, 3, 5])
+ CompareNonEmpty(self, [1, 2, 3], [1, 3, 1])
+ CompareNonEmpty(self, [1, 1, 3], [1, 3, 5])
+ CompareNonEmpty(self, [1, 2, 3], [1, 3, 5])
+ CompareNonEmpty(self, [7, 1, 3], [7, 3, 5])
+ CompareNonEmpty(self, [7, 2, 3], [7, 3, 1])
+ CompareNonEmpty(self, [7, 2, 3], [7, 3, 5])
+ CompareNonEmpty(self, [10, 64, 75], [10, 75, 30])
+ CompareNonEmpty(self, [5, 7, 2, 3], [5, 7, 3, 5])
def _testEmpty(self, dtype, adjoint_a, adjoint_b, use_static_shape):
- def compareEmpty(self, a_shape, b_shape):
+ def CompareEmpty(self, a_shape, b_shape):
self._compare(
np.zeros(a_shape).astype(dtype),
np.zeros(b_shape).astype(dtype), adjoint_a, adjoint_b,
use_static_shape)
- compareEmpty(self, [0, 3, 2], [0, 2, 4])
- compareEmpty(self, [3, 0, 2], [3, 2, 5])
- compareEmpty(self, [3, 3, 2], [3, 2, 0])
+ CompareEmpty(self, [0, 3, 2], [0, 2, 4])
+ CompareEmpty(self, [3, 0, 2], [3, 2, 5])
+ CompareEmpty(self, [3, 3, 2], [3, 2, 0])
def _GetBatchMatmulOpTest(dtype, adjoint_a, adjoint_b, use_static_shape):
+ @test_util.run_v1_only("b/120545219")
def Test(self):
np.random.seed(42)
self._testNonEmpty(dtype, adjoint_a, adjoint_b, use_static_shape)
@@ -154,17 +156,13 @@
y = y_in if not adjoint_b else y_in.reshape(y_t_shape)
epsilon = np.finfo(x.dtype).eps
delta = epsilon**(1.0 / 3.0)
+ def Loss(x, y):
+ z = math_ops.matmul(x, y, adjoint_a, adjoint_b)
+ return math_ops.reduce_sum(z)
with self.cached_session(use_gpu=True):
- inx = constant_op.constant(x)
- iny = constant_op.constant(y)
- z = math_ops.matmul(inx, iny, adjoint_a, adjoint_b)
- loss = math_ops.reduce_sum(z)
- ((x_jacob_t, x_jacob_n),
- (y_jacob_t, y_jacob_n)) = gradient_checker.compute_gradient(
- [inx, iny], [x.shape, y.shape],
- loss, [1],
- x_init_value=[x, y],
- delta=delta)
+ ((x_jacob_t, y_jacob_t),
+ (x_jacob_n, y_jacob_n)) = gradient_checker_v2.compute_gradient(
+ Loss, [x, y], delta=delta)
tol = 20 * delta
self.assertAllClose(x_jacob_t, x_jacob_n, rtol=tol, atol=tol)
self.assertAllClose(y_jacob_t, y_jacob_n, rtol=tol, atol=tol)
@@ -188,6 +186,7 @@
def _GetBatchMatmulGradientTest(dtype, adjoint_a, adjoint_b):
+ @test_util.run_v1_only("b/120545219")
def Test(self):
self._compare(1, 2, 3, 5, dtype, adjoint_a, adjoint_b)
self._compare(3, 4, 7, 10, dtype, adjoint_a, adjoint_b)
@@ -202,11 +201,12 @@
for adjoint_a_ in False, True:
for adjoint_b_ in False, True:
name = "%s_%s_%s" % (dtype_.__name__, adjoint_a_, adjoint_b_)
- for use_static_shape in True, False:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_static_shape_ in set([True, tf2.enabled()]):
setattr(BatchMatmulOpTest,
- "testBatchMatmulOp_" + name + ("_%s" % use_static_shape),
+ "testBatchMatmulOp_" + name + ("_%s" % use_static_shape_),
_GetBatchMatmulOpTest(dtype_, adjoint_a_, adjoint_b_,
- use_static_shape))
+ use_static_shape_))
if dtype_ is not np.int32:
setattr(BatchMatmulGradientTest, "testBatchMatmulGradient_" + name,
_GetBatchMatmulGradientTest(dtype_, adjoint_a_, adjoint_b_))
diff --git a/tensorflow/python/kernel_tests/batch_scatter_ops_test.py b/tensorflow/python/kernel_tests/batch_scatter_ops_test.py
index eefcdc5..f70fb93 100644
--- a/tensorflow/python/kernel_tests/batch_scatter_ops_test.py
+++ b/tensorflow/python/kernel_tests/batch_scatter_ops_test.py
@@ -22,6 +22,7 @@
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.ops import state_ops
from tensorflow.python.ops import variables
@@ -50,7 +51,8 @@
vtype,
itype,
repeat_indices=False,
- updates_are_scalar=False):
+ updates_are_scalar=False,
+ method=False):
np.random.seed(8)
with self.cached_session(use_gpu=False):
for indices_shape in (2,), (3, 7), (3, 4, 7):
@@ -71,7 +73,10 @@
# Scatter via tensorflow
ref = variables.Variable(old)
ref.initializer.run()
- tf_scatter(ref, indices, updates).eval()
+ if method:
+ ref.batch_scatter_update(ops.IndexedSlices(indices, updates))
+ else:
+ tf_scatter(ref, indices, updates).eval()
self.assertAllClose(ref.eval(), new)
@test_util.run_deprecated_v1
diff --git a/tensorflow/python/kernel_tests/boosted_trees/quantile_ops_test.py b/tensorflow/python/kernel_tests/boosted_trees/quantile_ops_test.py
index 390672f..2b9863f 100644
--- a/tensorflow/python/kernel_tests/boosted_trees/quantile_ops_test.py
+++ b/tensorflow/python/kernel_tests/boosted_trees/quantile_ops_test.py
@@ -82,7 +82,7 @@
self.max_elements = 1 << 16
self.num_quantiles = constant_op.constant(3, dtype=dtypes.int64)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBasicQuantileBucketsSingleResource(self):
with self.cached_session() as sess:
quantile_accumulator_handle = self.create_resource("floats", self.eps,
@@ -107,7 +107,7 @@
self.assertAllClose(self._feature_0_quantiles, quantiles[0].eval())
self.assertAllClose(self._feature_1_quantiles, quantiles[1].eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBasicQuantileBucketsMultipleResources(self):
with self.cached_session() as sess:
quantile_accumulator_handle_0 = self.create_resource("float_0", self.eps,
@@ -142,7 +142,7 @@
self.assertAllClose(self._feature_0_quantiles, quantiles[0].eval())
self.assertAllClose(self._feature_1_quantiles, quantiles[1].eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSaveRestoreAfterFlush(self):
save_dir = os.path.join(self.get_temp_dir(), "save_restore")
save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
@@ -175,7 +175,7 @@
self.assertAllClose(self._feature_0_boundaries, buckets[0].eval())
self.assertAllClose(self._feature_1_boundaries, buckets[1].eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSaveRestoreBeforeFlush(self):
save_dir = os.path.join(self.get_temp_dir(), "save_restore")
save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
diff --git a/tensorflow/python/kernel_tests/checkpoint_ops_test.py b/tensorflow/python/kernel_tests/checkpoint_ops_test.py
index b8c8c9e..dd5ac1f 100644
--- a/tensorflow/python/kernel_tests/checkpoint_ops_test.py
+++ b/tensorflow/python/kernel_tests/checkpoint_ops_test.py
@@ -125,7 +125,7 @@
save = saver.Saver([matrix])
with self.cached_session() as sess:
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.bundle_file = os.path.join(test.get_temp_dir(), 'bundle_checkpoint')
save.save(sess, self.bundle_file)
@@ -230,6 +230,7 @@
np.reshape(initializing_values, (num_rows, num_cols)),
self.evaluate(remapped_matrix))
+ @test_util.run_v1_only('b/120545219')
def test_load_and_remap_invalid_remapping(self):
"""Tests that errors are raised when an ID maps to multiple new IDs.
@@ -261,6 +262,7 @@
with self.cached_session(), self.assertRaises(errors.UnimplementedError):
self.evaluate(remapped_matrix)
+ @test_util.run_v1_only('b/120545219')
def test_load_and_remap_incorrect_initializing_values(self):
"""Tests that errors are raised with incorrect number of init values."""
remapped_matrix = gen_checkpoint_ops.load_and_remap_matrix(
@@ -312,7 +314,7 @@
with self.cached_session() as sess:
ckpt_path = os.path.join(test.get_temp_dir(), 'temp_ckpt')
save = saver.Saver([matrix])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
save.save(sess, ckpt_path)
num_rows, num_cols = np_value.shape
diff --git a/tensorflow/python/kernel_tests/cholesky_op_test.py b/tensorflow/python/kernel_tests/cholesky_op_test.py
index f394723..a08cfe9 100644
--- a/tensorflow/python/kernel_tests/cholesky_op_test.py
+++ b/tensorflow/python/kernel_tests/cholesky_op_test.py
@@ -155,6 +155,7 @@
np.array([[[1., 2., 3.], [3., 4., 5.]], [[1., 2., 3.], [3., 4., 5.]]
]))
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
tensor3 = constant_op.constant([1., 2.])
with self.assertRaises(ValueError):
@@ -233,6 +234,7 @@
self.runFiniteDifferences(
shapes, dtypes=(dtypes_lib.float64,), scalarTest=True)
+ @test_util.run_v1_only("b/120545219")
def testTwoBlockMatrixComplexFloat(self):
np.random.seed(0)
shapes = self.getShapes([2 * self._backprop_block_size + 1])
diff --git a/tensorflow/python/kernel_tests/cond_v2_test.py b/tensorflow/python/kernel_tests/cond_v2_test.py
index 1f4b37c..8fe3ba4 100644
--- a/tensorflow/python/kernel_tests/cond_v2_test.py
+++ b/tensorflow/python/kernel_tests/cond_v2_test.py
@@ -170,8 +170,8 @@
self.assertRegexpMatches(
cond2_op.get_attr("else_branch").name, r"foo_cond_1_false_\d*")
+ @test_util.run_v1_only("b/120545219")
def testDefunInCond(self):
- self.skipTest("b/117293122")
x = constant_op.constant(1.0, name="x")
y = constant_op.constant(2.0, name="y")
@@ -190,9 +190,8 @@
self._testCond(true_fn, false_fn, [x, y])
self._testCond(true_fn, false_fn, [y])
+ @test_util.run_deprecated_v1
def testNestedDefunInCond(self):
- self.skipTest("b/117284369")
-
x = constant_op.constant(1.0, name="x")
y = constant_op.constant(2.0, name="y")
@@ -216,9 +215,8 @@
self._testCond(true_fn, false_fn, [x, y])
self._testCond(true_fn, false_fn, [y])
+ @test_util.run_deprecated_v1
def testDoubleNestedDefunInCond(self):
- self.skipTest("b/117284369")
-
x = constant_op.constant(1.0, name="x")
y = constant_op.constant(2.0, name="y")
@@ -778,6 +776,26 @@
self.assertAllEqual(
self.evaluate(output_t), [-5, -4, -3, -2, -1, 0, 1, 4, 9, 16])
+ @test_util.run_deprecated_v1
+ def testForwardPassRewrite(self):
+ x = constant_op.constant(1.0, name="x")
+ output = cond_v2.cond_v2(constant_op.constant(True),
+ lambda: x * 2.0,
+ lambda: x)
+ if_op = output.op.inputs[0].op
+ self.assertEqual(if_op.type, "If")
+ # pylint: disable=g-deprecated-assert
+ self.assertEqual(len(if_op.outputs), 1)
+
+ gradients_impl.gradients(output, x)
+ # if_op should have been rewritten to output 2.0 intermediate.
+ self.assertEqual(len(if_op.outputs), 2)
+
+ gradients_impl.gradients(output, x)
+ # Computing the gradient again shouldn't rewrite if_op again.
+ self.assertEqual(len(if_op.outputs), 2)
+ # pylint: enable=g-deprecated-assert
+
class CondV2CollectionTest(test.TestCase):
diff --git a/tensorflow/python/kernel_tests/conditional_accumulator_test.py b/tensorflow/python/kernel_tests/conditional_accumulator_test.py
index 5847e46..32a2058 100644
--- a/tensorflow/python/kernel_tests/conditional_accumulator_test.py
+++ b/tensorflow/python/kernel_tests/conditional_accumulator_test.py
@@ -199,7 +199,7 @@
is_all_equal &= (val[i][j] == elems_ave[i][j])
self.assertTrue(is_all_equal)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAccumulatorWrongDynamicShape(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
@@ -321,7 +321,7 @@
shape=tensor_shape.TensorShape([1]),
reduction_type="Invalid")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAccumulatorInvalidTakeGrad(self):
with self.cached_session():
q = data_flow_ops.ConditionalAccumulator(
@@ -408,7 +408,7 @@
set_global_step_op = q.set_global_step(new_global_step)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
for _ in range(3):
set_global_step_op.run()
self.evaluate(inc_global_step)
@@ -435,7 +435,7 @@
if x >= ls) / sum(1 for x in local_steps
if x >= ls), val)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testParallelApplyGrad(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
@@ -461,7 +461,7 @@
self.assertEqual(val, sum(elems) / len(elems))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testParallelTakeGrad(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
@@ -494,7 +494,7 @@
self.assertItemsEqual(elems, results)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAccumulatorApplyAndBlockingTake(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
@@ -528,7 +528,7 @@
with self.assertRaisesOpError("was cancelled"):
self.evaluate(takeg_op)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAccumulatorCancel(self):
with self.cached_session() as sess:
q = data_flow_ops.ConditionalAccumulator(
diff --git a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
index f1efc5c..39ceb0d 100644
--- a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
+++ b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
@@ -33,6 +33,7 @@
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.eager import function as eager_function
+from tensorflow.python.eager import wrap_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
@@ -43,6 +44,7 @@
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 control_flow_util
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import gen_array_ops
@@ -129,6 +131,7 @@
@test_util.with_control_flow_v2
class ControlFlowTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testRefIdentity(self):
with self.cached_session():
v = variables.VariableV1(7)
@@ -138,10 +141,10 @@
v2 = control_flow_ops.with_dependencies([op], v)
self.assertTrue(isinstance(v2, ops.Tensor))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(9, self.evaluate(v2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRefEnter(self):
with self.cached_session():
v = variables.VariableV1(7)
@@ -152,10 +155,10 @@
op = state_ops.assign(enter_v, enter_nine)
v2 = control_flow_ops.with_dependencies([op], enter_v)
v3 = control_flow_ops.exit(v2)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(9, self.evaluate(v3))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRefSwitch(self):
with self.cached_session():
v = variables.VariableV1(7)
@@ -163,7 +166,7 @@
p = constant_op.constant(True)
v1 = control_flow_ops._SwitchRefOrTensor(v._ref(), p) # pylint: disable=protected-access
v2 = state_ops.assign(v1[1], 9)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(9, self.evaluate(v2))
def testEnterMulExit(self):
@@ -193,6 +196,7 @@
v, "frame2", is_constant=False)
self.assertEqual(enter_v_non_constant.shape, None)
+ @test_util.run_v1_only("b/120545219")
def testSwitchMergeIndexedSlices(self):
with self.cached_session():
values = constant_op.constant([1, 2, 3, 4, 5, 6])
@@ -202,11 +206,12 @@
switch_op = control_flow_ops.switch(data, pred)
merge_op = control_flow_ops.merge(switch_op)[0]
- val = merge_op.values.eval()
- ind = merge_op.indices.eval()
+ val = merge_op.values
+ ind = merge_op.indices
self.assertAllEqual(np.arange(1, 7), val)
self.assertAllEqual(np.arange(0, 12, 2), ind)
+ @test_util.run_v1_only("b/120545219")
def testSwitchDeadBranch(self):
with self.cached_session():
data = constant_op.constant([1, 2, 3, 4, 5, 6], name="data")
@@ -219,6 +224,7 @@
lambda e: "Retval[0] does not have value" in str(e)):
self.evaluate(dead_branch)
+ @test_util.run_v1_only("b/120545219")
def testSwitchMergeLess(self):
with self.cached_session():
data = constant_op.constant([1, 2, 3, 4, 5, 6], name="data")
@@ -231,6 +237,7 @@
result = self.evaluate(merge_op)
self.assertAllEqual(np.arange(1, 7), result)
+ @test_util.run_v1_only("b/120545219")
def testSwitchMergeAddIdentity(self):
with self.cached_session():
data = constant_op.constant([1, 2, 3, 4, 5, 6], name="data")
@@ -244,6 +251,7 @@
result = self.evaluate(merge_op)
self.assertAllEqual(np.array([x + 1 for x in [1, 2, 3, 4, 5, 6]]), result)
+ @test_util.run_v1_only("b/120545219")
def testSwitchMergeAddMul(self):
with self.cached_session():
data = constant_op.constant([1, 2, 3, 4, 5, 6], name="data")
@@ -258,6 +266,7 @@
result = self.evaluate(merge_op)
self.assertAllEqual(np.array([x * 5 for x in [1, 2, 3, 4, 5, 6]]), result)
+ @test_util.run_v1_only("b/120545219")
def testLoop_false(self):
with self.cached_session():
false = ops.convert_to_tensor(False)
@@ -302,6 +311,7 @@
result = self.evaluate(exit_i)
self.assertAllEqual(10, result)
+ @test_util.run_v1_only("b/120545219")
def testLoop_2(self):
with self.cached_session():
zero = constant_op.constant(0)
@@ -328,6 +338,7 @@
result = self.evaluate(exit_i)
self.assertAllEqual(10, result)
+ @test_util.run_v1_only("b/120545219")
def testDifferentFrame(self):
with self.cached_session():
data = array_ops.placeholder(dtypes.float32, shape=[])
@@ -362,6 +373,7 @@
lambda: math_ops.subtract(x, 1.))
self.assertEqual(b.shape, tensor_shape.scalar())
+ @test_util.run_v1_only("b/120545219")
def testFetchable(self):
with self.cached_session() as sess:
x = array_ops.placeholder(dtypes.float32)
@@ -378,6 +390,7 @@
sess.run(t, feed_dict={x: 3})
@test_util.disable_control_flow_v2("Not relevant")
+ @test_util.run_v1_only("b/120545219")
def testFeedable(self):
with self.cached_session() as sess:
c = constant_op.constant(2)
@@ -395,6 +408,7 @@
with self.assertRaisesRegexp(ValueError, "may not be fed"):
sess.run(r, feed_dict={t: 3})
+ @test_util.run_v1_only("b/120545219")
def testCondIndexedSlices(self):
with self.cached_session():
values = constant_op.constant(10)
@@ -405,11 +419,12 @@
fn2 = lambda: ops.IndexedSlices(math_ops.subtract(x.values, 1), indices)
r = control_flow_ops.cond(pred, fn1, fn2)
- val = r.values.eval()
- ind = r.indices.eval()
+ val = r.values
+ ind = r.indices
self.assertAllEqual(11, val)
self.assertAllEqual(0, ind)
+ @test_util.run_v1_only("b/120545219")
def testCondSparseTensor(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -423,15 +438,16 @@
fn2 = lambda: sparse_tensor.SparseTensor(
indices, x.values - 1, dense_shape=shape)
r = control_flow_ops.cond(pred, fn1, fn2)
- self.assertAllEqual([3.0, 5.0], r.values.eval())
- self.assertAllEqual([[1], [4]], r.indices.eval())
+ self.assertAllEqual([3.0, 5.0], r.values)
+ self.assertAllEqual([[1], [4]], r.indices)
self.assertAllEqual(r.values.get_shape(), (2,))
+ @test_util.run_v1_only("b/120545219")
def testCondResource(self):
with self.cached_session():
rv = resource_variable_ops.ResourceVariable(True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
t = ops.convert_to_tensor(1.0)
def case():
@@ -439,8 +455,10 @@
with ops.control_dependencies([assign]):
return array_ops.identity(t)
- self.assertEqual(1.0, control_flow_ops.cond(rv, case, lambda: t).eval())
+ self.assertEqual(
+ 1.0, self.evaluate(control_flow_ops.cond(rv, case, lambda: t)))
+ @test_util.run_v1_only("b/120545219")
def testCondWithTensorArrayGrad(self):
with self.cached_session() as sess:
with ops.device(test.gpu_device_name()):
@@ -455,6 +473,7 @@
self.assertAllEqual(sess.run(g, {pred: False}), [0.0, 0.0, 0.0])
@test_util.disable_control_flow_v2("b/113293074")
+ @test_util.run_v1_only("b/120545219")
def testCondIndexedSlicesDifferentTypes(self):
with self.cached_session():
values = constant_op.constant(10)
@@ -466,12 +485,13 @@
fn2 = lambda: ops.IndexedSlices(math_ops.subtract(x.values, 1), i_64)
r = control_flow_ops.cond(pred, fn1, fn2)
- val = r.values.eval()
- ind = r.indices.eval()
+ val = r.values
+ ind = r.indices
self.assertAllEqual(11, val)
self.assertAllEqual(0, ind)
self.assertTrue(ind.dtype == np.int64)
+ @test_util.run_v1_only("b/120545219")
def testCondColocation(self):
with self.session(use_gpu=True):
with ops.device("/cpu:0"):
@@ -547,8 +567,8 @@
if not context.executing_eagerly():
with self.cached_session():
- variables.global_variables_initializer().run()
- result = f().eval()
+ self.evaluate(variables.global_variables_initializer())
+ result = self.evaluate(f())
self.assertEqual(True, result)
# Only second cond result was fetched, so v1 assign shouldn't run.
self.assertEqual(7, self.evaluate(v1))
@@ -576,6 +596,7 @@
alive, count = body(i)
self.assertAllEqual(4, self.evaluate(count))
+ @test_util.run_v1_only("b/120545219")
def testCond_6(self):
with self.cached_session():
v1 = variables.Variable([7])
@@ -586,7 +607,7 @@
fn2 = lambda: v1
r = control_flow_ops.cond(pred, fn1, fn2)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
result = self.evaluate(r)
self.assertAllEqual(np.array([7]), result)
@@ -671,6 +692,7 @@
test_result = self.evaluate(r)
self.assertDictEqual({"a": {"c": 210}, "b": {"d": 210}}, test_result)
+ @test_util.run_v1_only("b/120545219")
def testCheckNestedOutputStruct(self):
with self.cached_session() as sess:
x = constant_op.constant(10)
@@ -681,7 +703,8 @@
v1_msg = "The two structures don't have the same nested structure"
v2_msg = "Outputs of true_fn and false_fn must have the same structure"
with self.assertRaisesRegexp(
- ValueError, v2_msg if control_flow_ops.ENABLE_COND_V2 else v1_msg):
+ ValueError,
+ v2_msg if control_flow_util.ENABLE_CONTROL_FLOW_V2 else v1_msg):
r = control_flow_ops.cond(pred, fn1, fn2)
self.evaluate(r)
@@ -701,7 +724,7 @@
self.assertAllEqual([2.0], self.evaluate(r))
@test_util.disable_control_flow_v2("b/79881896 (control deps)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCondWithControl(self):
with self.cached_session():
control_holder = array_ops.placeholder(dtypes.float32, shape=())
@@ -717,6 +740,7 @@
lambda: constant_op.constant(1))
self.assertEqual(5, self.evaluate(r))
+ @test_util.run_v1_only("b/120545219")
def testUninitializedRefIdentity(self):
with self.cached_session() as sess:
v = gen_state_ops.variable(
@@ -771,6 +795,7 @@
r = control_flow_ops.cond(pred, fn1, fn2)
self.evaluate(r)
+ @test_util.run_v1_only("b/120545219")
def testCondGrad_1(self):
with self.cached_session():
x = constant_op.constant(10.0, name="x")
@@ -831,20 +856,21 @@
with ops.device("/cpu:1"):
grad = gradients_impl.gradients(z, x)[0]
- self.assertEqual(sess.run(grad, {pred: True, x: 1.0, y: 2.0}), 4.0)
- self.assertEqual(sess.run(grad, {pred: False, x: 1.0, y: 2.0}), 0.0)
-
with ops.device("/cpu:0"):
grad_grad = gradients_impl.gradients(grad, x)[0]
+ self.assertEqual(sess.run(grad, {pred: True, x: 1.0, y: 2.0}), 4.0)
+ self.assertEqual(sess.run(grad, {pred: False, x: 1.0, y: 2.0}), 0.0)
+
# v1 control flow gets None second derivative for some reason.
- if not control_flow_ops.ENABLE_COND_V2:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
self.assertIsNone(grad_grad)
return
self.assertEqual(sess.run(grad_grad, {pred: True, x: 1.0, y: 2.0}), 0.0)
self.assertEqual(sess.run(grad_grad, {pred: False, x: 1.0, y: 2.0}), 0.0)
+ @test_util.run_v1_only("b/120545219")
def testNestedCond_Simple(self):
with self.cached_session():
x = constant_op.constant(0., name="X")
@@ -861,7 +887,7 @@
self.assertEqual(1.0, self.evaluate(result))
@test_util.disable_control_flow_v2("b/113327884")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCondGrad_Gather(self):
with self.cached_session() as sess:
v1 = variables.Variable([1.0, 42.0])
@@ -871,7 +897,7 @@
fn2 = lambda: array_ops.gather(v1, [1, 1])
r = control_flow_ops.cond(pred, fn1, fn2)
grad = gradients_impl.gradients(r, [v1])[0]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Should just be [1, 1], but possibly a sparse representation
gv, gi = sess.run([grad.values, grad.indices], feed_dict={c: 1})
dense_gv = [
@@ -885,6 +911,7 @@
]
self.assertAllEqual(dense_gv, [0.0, 2.0])
+ @test_util.run_v1_only("b/120545219")
def testCondPredicateTensor(self):
"""Regression test for lowering predicate from non-first output of an op."""
@@ -895,9 +922,8 @@
r = control_flow_ops.cond(foo()[1], lambda: 1.0, lambda: 2.0)
self.assertEqual(self.evaluate(r), 1.0)
- # TODO(b/117945658): reenable
@test_util.run_in_graph_and_eager_modes
- def DISABLED_testCondAutoControlDeps(self):
+ def testCondAutoControlDeps(self):
def branch_fn():
logging_ops.print_v2("A")
@@ -917,16 +943,16 @@
if not context.executing_eagerly():
with self.cached_session():
with self.captureWritesToStream(sys.stderr) as printed:
- self.assertEqual(build_cond().eval(), 10)
+ self.assertEqual(self.evaluate(build_cond()), 10)
self.assertEqual(printed.contents(), "C\n")
with self.captureWritesToStream(sys.stderr) as printed:
- self.assertEqual(build_nested_cond().eval(), 10)
+ self.assertEqual(self.evaluate(build_nested_cond()), 10)
self.assertEqual(printed.contents(), "C\n")
# In defuns, all prints should execute in program order.
# This doesn't work with legacy control flow.
- if control_flow_ops.ENABLE_COND_V2:
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2:
@eager_function.defun
def cond():
@@ -944,9 +970,28 @@
self.assertEqual(self.evaluate(nested_cond()), 10)
self.assertEqual(printed.contents(), "A\nB\nC\n")
- # TODO(b/117945658): reenable
+ # wrap_function should prune.
+ def pruned_cond():
+ return build_cond()
+ pruned_cond = wrap_function.wrap_function(pruned_cond, [])
+
+ with self.captureWritesToStream(sys.stderr) as printed:
+ self.assertEqual(self.evaluate(pruned_cond()), 10)
+ self.assertEqual(printed.contents(), "C\n")
+
+ def pruned_nested_cond():
+ return build_nested_cond()
+ pruned_nested_cond = wrap_function.wrap_function(pruned_nested_cond, [])
+
+ with self.captureWritesToStream(sys.stderr) as printed:
+ self.assertEqual(self.evaluate(pruned_nested_cond()), 10)
+ self.assertEqual(printed.contents(), "C\n")
+
@test_util.run_in_graph_and_eager_modes
- def DISABLED_testWhileAutoControlDeps(self):
+ def testWhileAutoControlDeps(self):
+ # Legacy while_loop fails this test because it produces deprecation notices
+ # in stderr.
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2: return
def cond(i, unused_x):
logging_ops.print_v2("A")
@@ -965,40 +1010,56 @@
def build_nested_while():
return control_flow_ops.cond(
- constant_op.constant(True), build_while, lambda: (0, 0))
+ constant_op.constant(True), build_while, lambda: [0, 0])
# In v1 graph mode, pruning should make only "D" print.
if not context.executing_eagerly():
with self.cached_session():
with self.captureWritesToStream(sys.stderr) as printed:
- self.assertEqual(build_while()[0].eval(), 2)
+ self.assertEqual(self.evaluate(build_while()[0]), 2)
self.assertEqual(printed.contents(), "D\nD\n")
with self.captureWritesToStream(sys.stderr) as printed:
- self.assertEqual(build_nested_while()[0].eval(), 2)
+ self.assertEqual(self.evaluate(build_nested_while()[0]), 2)
self.assertEqual(printed.contents(), "D\nD\n")
# In defuns, all prints should execute in program order.
- # This doesn't work with legacy control flow.
- if control_flow_ops.ENABLE_WHILE_V2:
+ @eager_function.defun
+ def while_loop():
+ return build_while()[0]
- @eager_function.defun
- def while_loop():
- return build_while()[0]
+ with self.captureWritesToStream(sys.stderr) as printed:
+ self.assertEqual(self.evaluate(while_loop()), 2)
+ self.assertEqual(printed.contents(), "A\nB\nC\nD\nA\nB\nC\nD\nA\n")
+ @eager_function.defun
+ def nested_while_loop():
+ return build_nested_while()[0]
+
+ # TODO(b/117840611): calling nested_while_loop fails in eager
+ if not context.executing_eagerly():
with self.captureWritesToStream(sys.stderr) as printed:
- self.assertEqual(self.evaluate(while_loop()), 2)
+ self.assertEqual(self.evaluate(nested_while_loop()), 2)
self.assertEqual(printed.contents(), "A\nB\nC\nD\nA\nB\nC\nD\nA\n")
- @eager_function.defun
- def nested_while_loop():
- return build_nested_while()[0]
+ # wrap_function should prune.
+ def pruned_while():
+ return build_while()[0]
+ pruned_while = wrap_function.wrap_function(pruned_while, [])
- # TODO(b/117840611): calling nested_while_loop fails in eager
- if not context.executing_eagerly():
- with self.captureWritesToStream(sys.stderr) as printed:
- self.assertEqual(self.evaluate(nested_while_loop()), 2)
- self.assertEqual(printed.contents(), "A\nB\nC\nD\nA\nB\nC\nD\nA\n")
+ with self.captureWritesToStream(sys.stderr) as printed:
+ self.assertEqual(self.evaluate(pruned_while()), 2)
+ self.assertEqual(printed.contents(), "D\nD\n")
+
+ def pruned_nested_while():
+ return build_nested_while()[0]
+ pruned_nested_while = wrap_function.wrap_function(pruned_nested_while, [])
+
+ # TODO(b/117840611): calling nested_while_loop fails in eager
+ if not context.executing_eagerly():
+ with self.captureWritesToStream(sys.stderr) as printed:
+ self.assertEqual(self.evaluate(pruned_nested_while()), 2)
+ self.assertEqual(printed.contents(), "D\nD\n")
# Microbenchmark: 256,000 iterations/s.
@test_util.disable_control_flow_v2("b/116630618 (Times out)")
@@ -1011,6 +1072,7 @@
self.assertEqual(10000, self.evaluate(r))
@test_util.disable_control_flow_v2("b/79881896 (control deps)")
+ @test_util.run_v1_only("b/120545219")
def testWhileExternalControlDependencies(self):
with self.cached_session():
v = variables.Variable(0.0)
@@ -1023,10 +1085,11 @@
result = control_flow_ops.while_loop(cond=lambda i: i < 2,
body=body_fn, loop_vars=[1])
- self.assertAllEqual(result.eval(), 2)
- self.assertAllEqual(v.eval(), 1.0)
+ self.assertAllEqual(result, 2)
+ self.assertAllEqual(v.read_value(), 1.0)
@test_util.disable_control_flow_v2("b/79881896 (control deps)")
+ @test_util.run_v1_only("b/120545219")
def testWhileExternalControlDependenciesNoInput(self):
with self.cached_session():
v = variables.Variable(0.0)
@@ -1040,10 +1103,10 @@
result = control_flow_ops.while_loop(cond=lambda i: i < 5,
body=body_fn, loop_vars=[0])
self.evaluate(result)
- self.assertAllEqual(v.eval(), 1.0)
+ self.assertAllEqual(self.evaluate(v), 1.0)
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileWithRefs_1(self):
with self.cached_session() as sess:
x = variables.VariableV1(0)._ref() # pylint: disable=protected-access
@@ -1058,7 +1121,7 @@
r = control_flow_ops.while_loop(c, b, [i, x], parallel_iterations=5)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(r[0].dtype, dtypes.int32)
self.assertEqual(r[1].dtype, dtypes.int32_ref)
@@ -1080,6 +1143,7 @@
r = isum(s, maximum_iterations=3)
self.assertAllEqual([1 + 3, 2 + 3, 3 + 3, 4 + 3, 5 + 3], self.evaluate(r))
+ @test_util.run_v1_only("b/120545219")
def testWhileWithMaximumIterationsAndSingleArgument(self):
with self.cached_session():
r = control_flow_ops.while_loop(
@@ -1087,6 +1151,7 @@
self.assertEqual(1, self.evaluate(r))
@test_util.disable_control_flow_v2("b/115776323 (max_iters)")
+ @test_util.run_v1_only("b/120545219")
def testSingleNestedMaximumIterationsWhileLoopGradientInXLAContext(self):
v = constant_op.constant(1.0)
@@ -1112,6 +1177,7 @@
# Should execute without issue.
self.assertEqual(3, self.evaluate(loop_execute))
+ @test_util.run_v1_only("b/120545219")
def testInvalidMaximumIterationsWhileLoopGradientInXLAContext(self):
v = constant_op.constant(1.0)
@@ -1133,7 +1199,7 @@
gs = gradients_impl.gradients(loop_no_xla, v)
self.evaluate(gs) # This should execute without error.
- if control_flow_ops.ENABLE_WHILE_V2:
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2:
xla_context = control_flow_ops.XLAControlFlowContext()
xla_context.Enter()
with self.assertRaisesRegexp(
@@ -1172,6 +1238,7 @@
r"context '.*' \(currently defined in '.*'\)"):
_ = gradients_impl.gradients(loop_with_maxiter, v)
+ @test_util.run_v1_only("b/120545219")
def testInvalidMaximumIterationsFromSiblingContextWhileLoopInXLAContext(self):
v = constant_op.constant(1.0)
@@ -1190,7 +1257,7 @@
lambda i, x: (i + 1, v * x), (0, 1.0),
maximum_iterations=max_iter_holder[0])
- if control_flow_ops.ENABLE_WHILE_V2:
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2:
xla_context = control_flow_ops.XLAControlFlowContext()
xla_context.Enter()
with self.assertRaisesRegexp(
@@ -1215,6 +1282,7 @@
_ = gradients_impl.gradients(loop, v)
@test_util.disable_control_flow_v2("b/118457764")
+ @test_util.run_v1_only("b/120545219")
def testNestedWhileLoopWithMaxItersFromOuterContextInXLAContext(self):
v = constant_op.constant(1.0)
@@ -1301,7 +1369,7 @@
d = ops.convert_to_tensor(100)
r = control_flow_ops.while_loop(lambda i, m, c, o: math_ops.less(i, d),
compute, [i, m, c, o])
- result = r[3].eval()
+ result = r[3]
self.assertAllEqual(10100, result)
@test_util.run_deprecated_v1
@@ -1323,9 +1391,10 @@
s = array_ops.size(x)
r = control_flow_ops.while_loop(lambda i, m, c, o: math_ops.less(i, s),
compute, [i, m, c, o])
- result = r[3].eval()
+ result = r[3]
self.assertAllEqual(42, result)
+ @test_util.run_v1_only("b/120545219")
def testWhile_5(self):
with self.cached_session():
@@ -1347,10 +1416,11 @@
tensor_shape.unknown_shape(),
tensor_shape.unknown_shape()
])
- result = r[2].eval()
+ result = r[2]
self.assertAllEqual(np.array([0, 1, 2, 3, 4, 5, 6]), result)
@test_util.disable_control_flow_v2("b/116338794 (buffer_reuse)")
+ @test_util.run_v1_only("b/120545219")
def testBufferForwarding(self):
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
@@ -1435,6 +1505,7 @@
r = control_flow_ops.while_loop(c, b, [n], parallel_iterations=20)
self.assertEqual([10000], self.evaluate(r))
+ @test_util.run_v1_only("b/120545219")
def testWhileShapeInference(self):
with self.cached_session():
i = constant_op.constant(0)
@@ -1461,6 +1532,7 @@
r = control_flow_ops.while_loop(c, b, [i, m])
@test_util.disable_control_flow_v2("b/116328420 (SparseTensor)")
+ @test_util.run_v1_only("b/120545219")
def testWhileShapeInferenceSparseTensor(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -1480,12 +1552,12 @@
]
_, r = control_flow_ops.while_loop(c, b, [i, x])
- self.assertEqual(r.dense_shape.get_shape()[0].value, 1)
+ self.assertEqual(r.dense_shape.get_shape()[0], 1)
_, r = control_flow_ops.while_loop(
c, b, [i, x],
[i.get_shape(), tensor_shape.TensorShape([None])])
- self.assertTrue(r.dense_shape.get_shape()[0].value is None)
+ self.assertEqual(r.dense_shape.get_shape().as_list(), [None])
with self.assertRaisesRegexp(ValueError, "is not compatible with"):
_, r = control_flow_ops.while_loop(
@@ -1493,7 +1565,7 @@
[i.get_shape(), tensor_shape.TensorShape([5])])
@test_util.disable_control_flow_v2("b/116282023 (IndexedSlices)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileShapeInferenceIndexedSlices(self):
with self.cached_session():
values = constant_op.constant([[2.0, 4.0], [3.0, 5.0]], name="values")
@@ -1512,15 +1584,14 @@
]
_, r = control_flow_ops.while_loop(c, b, [i, x])
- self.assertEqual(r.dense_shape.get_shape()[0].value, 2)
+ self.assertEqual(r.dense_shape.get_shape()[0], 2)
self.assertEqual(r.values.get_shape(), tensor_shape.TensorShape([2, 2]))
_, r = control_flow_ops.while_loop(
c, b, [i, x],
[i.get_shape(), tensor_shape.TensorShape([None, 2])])
- self.assertEqual(r.dense_shape.get_shape()[0].value, 2)
- self.assertTrue(r.values.get_shape()[0].value is None)
- self.assertEqual(r.values.get_shape()[1].value, 2)
+ self.assertEqual(r.dense_shape.get_shape()[0], 2)
+ self.assertEqual(r.values.get_shape().as_list(), [None, 2])
with self.assertRaisesRegexp(ValueError, "is not compatible with"):
_, r = control_flow_ops.while_loop(
@@ -1584,6 +1655,7 @@
self._testNestedWhile_2(use_gpu=False)
self._testNestedWhile_2(use_gpu=True)
+ @test_util.run_v1_only("b/120545219")
def testWhileWithControl_1(self):
with self.cached_session():
n = constant_op.constant(0)
@@ -1598,7 +1670,7 @@
res = control_flow_ops.while_loop(
condition, body, [n, r], parallel_iterations=1)
- self.assertAllEqual(12, res[1].eval())
+ self.assertAllEqual(12, res[1])
@test_util.run_deprecated_v1
def testWhileWithControl_2(self):
@@ -1615,6 +1687,7 @@
condition, body, [r], parallel_iterations=1)
self.assertAllEqual(12, self.evaluate(res))
+ @test_util.run_v1_only("b/120545219")
def testWhileWithControl_3(self):
with self.cached_session() as sess:
b = array_ops.placeholder(dtypes.bool)
@@ -1624,6 +1697,7 @@
r = control_flow_ops.while_loop(lambda x: x < 10, lambda x: x + c, [x0])
self.assertEqual(10, sess.run(r, {b: True}))
+ @test_util.run_v1_only("b/120545219")
def testWhileWithControl_4(self):
with self.cached_session() as sess:
b = array_ops.placeholder(dtypes.bool)
@@ -1635,6 +1709,7 @@
self.assertEqual(10, sess.run(r, {b: True}))
@test_util.disable_control_flow_v2("b/79881896 (control_deps)")
+ @test_util.run_v1_only("b/120545219")
def testWhileWithControl_5(self):
with self.cached_session() as sess:
b = array_ops.placeholder(dtypes.bool)
@@ -1663,6 +1738,7 @@
self.assertEqual(0, self.evaluate(loop))
@test_util.disable_control_flow_v2("b/113324949 (ref vars)")
+ @test_util.run_v1_only("b/120545219")
def testWhileCondWithControl_1(self):
with self.cached_session():
v = variable_scope.get_variable(
@@ -1681,11 +1757,12 @@
return i + 1
r = control_flow_ops.while_loop(loop_condition, loop_body, (i0,))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(4, self.evaluate(r))
self.assertAllClose(65536.0, self.evaluate(v))
@test_util.disable_control_flow_v2("b/113324949 (ref vars)")
+ @test_util.run_v1_only("b/120545219")
def testWhileCondExitControl(self):
with self.cached_session():
@@ -1706,7 +1783,7 @@
r = control_flow_ops.cond(
constant_op.constant(False), lambda: constant_op.constant(1.0),
false_branch)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(6.0, self.evaluate(r))
self.assertEqual(99, self.evaluate(v))
@@ -1817,15 +1894,15 @@
with ops.device("/cpu:1"):
grad = gradients_impl.gradients(z, x_init)[0]
+ with ops.device("/cpu:0"):
+ grad_grad = gradients_impl.gradients(grad, x_init)[0]
+
self.assertEqual(sess.run(grad, {pred: True}), 8.0)
self.assertEqual(sess.run(grad, {pred: False}), 0.0)
- if not control_flow_ops.ENABLE_WHILE_V2:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
return
- with ops.device("/cpu:0"):
- grad_grad = gradients_impl.gradients(grad, x_init)[0]
-
self.assertEqual(sess.run(grad_grad, {pred: True}), 0.0)
self.assertEqual(sess.run(grad_grad, {pred: False}), 0.0)
@@ -1849,12 +1926,13 @@
r = control_flow_ops.while_loop(
loop_iterator, loop_body, [n], parallel_iterations=1)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(3, self.evaluate(r))
result = self.evaluate(select)
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result)
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
+ @test_util.run_v1_only("b/120545219")
def testWhileUpdateVariable_2(self):
with self.cached_session():
select1 = variables.Variable([3.0, 4.0, 5.0])
@@ -1874,7 +1952,7 @@
r = control_flow_ops.while_loop(
loop_iterator, loop_body, [n], parallel_iterations=1)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(3, self.evaluate(r))
result1 = self.evaluate(select1)
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result1)
@@ -1882,7 +1960,7 @@
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result2)
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileUpdateVariable_3(self):
with self.cached_session():
select = variables.Variable([3.0, 4.0, 5.0])
@@ -1900,17 +1978,17 @@
loop_iterator,
loop_body, [n, array_ops.identity(select)],
parallel_iterations=1)
- variables.global_variables_initializer().run()
- result = r[1].eval()
+ self.evaluate(variables.global_variables_initializer())
+ result = r[1]
self.assertAllClose(np.array([10.0, 10.0, 10.0]), result)
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileUpdateVariable_4(self):
with self.cached_session():
var_a = variables.Variable(0, name="a")
var_b = variables.Variable(0, name="b")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
c = constant_op.constant(0, name="c")
asn1 = state_ops.assign_add(var_a, 1, name="a_add")
@@ -1934,13 +2012,13 @@
self.assertEqual(10, self.evaluate(var_b))
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileUpdateVariable_5(self):
with self.cached_session():
# Create some variables.
var_a = variables.Variable(0, name="a")
var_b = variables.Variable(0, name="b")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Change condition to check var_b
def pred(_):
@@ -1965,13 +2043,14 @@
self.assertEqual(10, self.evaluate(var_b))
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
+ @test_util.run_v1_only("b/120545219")
def testWhileUpdateVariable_6(self):
with self.cached_session():
# Create some variables.
var_a = variables.Variable(0, name="a")
var_b = variables.Variable(0, name="b")
c = constant_op.constant(0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Loop condition
def pred(i):
@@ -1994,6 +2073,7 @@
self.assertEqual(55, self.evaluate(var_b))
self.assertEqual(10, self.evaluate(var_a))
+ @test_util.run_v1_only("b/120545219")
def testWhileQueue_1(self):
with self.cached_session():
q = data_flow_ops.FIFOQueue(-1, dtypes.int32)
@@ -2010,9 +2090,21 @@
r = control_flow_ops.while_loop(c, b, [i], parallel_iterations=1)
self.assertEqual([10], self.evaluate(r))
for i in xrange(10):
- self.assertEqual([i], q.dequeue().eval())
+ self.assertEqual([i], self.evaluate(q.dequeue()))
+
+ @test_util.run_v1_only("b/120545219")
+ def testWhileTimeOut(self):
+ run_options = config_pb2.RunOptions(timeout_in_ms=1)
+ with self.cached_session() as sess:
+ n = constant_op.constant(0)
+ c = lambda x: True
+ b = lambda x: math_ops.add(x, 1)
+ r = control_flow_ops.while_loop(c, b, [n])
+ with self.assertRaises(errors_impl.DeadlineExceededError):
+ sess.run(r, options=run_options)
@test_util.disable_control_flow_v2("b/117119329 (stack)")
+ @test_util.run_v1_only("b/120545219")
def testWhileStack_1(self):
with self.cached_session():
s = gen_data_flow_ops.stack_v2(-1, dtypes.int32, stack_name="foo")
@@ -2082,10 +2174,12 @@
self.assertAllClose(1024.0, self.evaluate(r))
@test_util.disable_control_flow_v2("b/116351701 (colocation)")
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_ColocateGradients(self):
self._testWhileGrad_ColocateGradients(colocate=False)
self._testWhileGrad_ColocateGradients(colocate=True)
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_Square(self):
with self.cached_session():
v = constant_op.constant(2.0, name="v")
@@ -2097,6 +2191,7 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(1024.0, self.evaluate(r))
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_Shape(self):
with self.cached_session():
x = array_ops.placeholder(dtypes.float32, shape=[None])
@@ -2127,6 +2222,7 @@
r = gradients_impl.gradients([r, y], x)[0]
self.assertAllClose([2.0, 4.0], sess.run(r, feed_dict={x: [1.0, 2.0]}))
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_MultipleUses(self):
with self.cached_session():
v = constant_op.constant(2.0, name="v")
@@ -2138,6 +2234,7 @@
r = gradients_impl.gradients(r, v)[0]
self.assertEqual(524288.0, self.evaluate(r))
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_LoopAdd(self):
with self.cached_session():
v = constant_op.constant(2.0, name="v")
@@ -2201,6 +2298,7 @@
def testNestedWhileCondWhileGradGpu(self):
self._testNestedWhileCondWhileGrad(use_gpu=True)
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_Variable(self):
with self.cached_session():
a = variables.Variable(3.0)
@@ -2210,8 +2308,8 @@
r = control_flow_ops.while_loop(c, b, [v], parallel_iterations=1)
r = gradients_impl.gradients(r, a)
- variables.global_variables_initializer().run()
- self.assertAllClose(216.0, r[0].eval())
+ self.evaluate(variables.global_variables_initializer())
+ self.assertAllClose(216.0, r[0])
@test_util.run_deprecated_v1
def testWhileGrad_ResourceVariable(self):
@@ -2223,9 +2321,10 @@
r = control_flow_ops.while_loop(c, b, [v], parallel_iterations=1)
g = gradients_impl.gradients(r, a)
- variables.global_variables_initializer().run()
- self.assertAllClose(216.0, g[0].eval())
+ self.evaluate(variables.global_variables_initializer())
+ self.assertAllClose(216.0, g[0])
+ @test_util.run_v1_only("b/120545219")
def testWhileGradInCond(self):
with self.cached_session():
@@ -2243,7 +2342,7 @@
self.assertAllClose(9.0, r.eval(feed_dict={x: 1.0}))
@test_util.disable_control_flow_v2("b/116340060")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testGradInWhileWrtInitialLoopVal(self):
with self.cached_session():
x = array_ops.placeholder(dtypes.float32, shape=(), name="x")
@@ -2261,6 +2360,7 @@
"loop invariants or wrt the input parameters to the loop body."):
control_flow_ops.while_loop(lambda i, x: i < 3, body, [0, y])
+ @test_util.run_v1_only("b/120545219")
def testWhileGradInWhile(self):
with self.cached_session():
n = ops.convert_to_tensor(1.0, name="n")
@@ -2277,6 +2377,7 @@
[tensor_shape.unknown_shape()])
self.assertAllClose(9.0, r.eval(feed_dict={x: 1.0}))
+ @test_util.run_v1_only("b/120545219")
def testCondGradInNestedWhiles(self):
def outer_body(i, x):
@@ -2295,6 +2396,49 @@
self.assertEqual(i_val, 3)
self.assertAllClose(x_val, 1.0)
+ def testNestedResourceAccess(self):
+ var = resource_variable_ops.ResourceVariable(constant_op.constant(3.0))
+
+ @eager_function.defun
+ def test_fn():
+ x = constant_op.constant(0.0)
+ r = control_flow_ops.while_loop(
+ # Outer loop condition
+ lambda i, y: i < 2,
+ # Outer loop body
+ lambda i, y: (i + 1, y + control_flow_ops.cond(
+ constant_op.constant(True),
+ # True branch
+ lambda: control_flow_ops.while_loop(
+ # Inner loop condition
+ lambda j, z: j < 3,
+ # Inner loop body
+ lambda j, z: (j + 1, z + math_ops.square(var)),
+ # Inner initial loop value
+ [0, y])[1],
+ # False branch
+ lambda: (0.0))),
+ # Outer initial loop value
+ [0, x])[1]
+
+ grad = gradients_impl.gradients(r, x)[0]
+ return r, grad
+
+ self.evaluate(variables.global_variables_initializer())
+ r, grad = self.evaluate(test_fn())
+ # 2 * 3 * 3^2
+ self.assertEqual(r, 81.0)
+ # v1 control flow gets the wrong answer!!!
+ # Gradient computation:
+ # f(x) = x + 3^2
+ # inner_loop(x) = f(f(f(x))) = x + 3*3^2 = x + 27
+ # g(x) = x + inner_loop(x) = 2x + 27
+ # outer_loop(x) = g(g(x)) = 4x + 81
+ # outer_loop'(x) = 4
+ # Note that v1 control flow gets 4.0 as well if the cond is removed.
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2:
+ self.assertEqual(grad, 4.0)
+
def testWhile_NestedInput(self):
with self.cached_session() as sess:
named = collections.namedtuple("named", ("a", "b"))
@@ -2322,6 +2466,7 @@
self.assertEqual([100.0, 1.0, 102.0, 3.0, 4.0 + 100 * 2.0],
self.evaluate(r_flattened))
+ @test_util.run_v1_only("b/120545219")
def testWhile_NestedBadArityFails(self):
with self.cached_session():
named = collections.namedtuple("named", ("a", "b"))
@@ -2338,6 +2483,7 @@
with self.assertRaisesRegexp(ValueError, "the same number of elements"):
control_flow_ops.while_loop(c, b, loop_vars)
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_ys_xs(self):
with self.cached_session():
x = constant_op.constant(3.0, name="x")
@@ -2353,13 +2499,13 @@
rx, ry = control_flow_ops.while_loop(c, b, [x, y], parallel_iterations=1)
r = gradients_impl.gradients([rx, ry], x)
- self.assertAllClose(304.0, r[0].eval())
+ self.assertAllClose(304.0, r[0])
r = gradients_impl.gradients([rx, ry], y)
- self.assertAllClose(124.0, r[0].eval())
+ self.assertAllClose(124.0, r[0])
r = gradients_impl.gradients([rx], x)
- self.assertAllClose(295.0, r[0].eval())
+ self.assertAllClose(295.0, r[0])
r = gradients_impl.gradients([rx], y)
- self.assertAllClose(120.0, r[0].eval())
+ self.assertAllClose(120.0, r[0])
@test_util.run_deprecated_v1
def testWhileGrad_Dependency(self):
@@ -2377,11 +2523,12 @@
ri, rx = control_flow_ops.while_loop(c, b, [i, x], parallel_iterations=1)
r = gradients_impl.gradients([ri, rx], x)
- self.assertAllClose(1024.0, r[0].eval())
+ self.assertAllClose(1024.0, r[0])
r = gradients_impl.gradients([rx], x)
- self.assertAllClose(1024.0, r[0].eval())
+ self.assertAllClose(1024.0, r[0])
@test_util.disable_control_flow_v2("b/116355153 (back_prop flag)")
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_NoGradient(self):
with self.cached_session():
v = constant_op.constant(2.0, name="v")
@@ -2390,9 +2537,10 @@
r = control_flow_ops.while_loop(c, b, [v], back_prop=False)
r = math_ops.add(r, v)
r = gradients_impl.gradients(r, v)
- self.assertAllClose(1.0, r[0].eval())
+ self.assertAllClose(1.0, r[0])
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_NoDependency(self):
with self.cached_session() as sess:
variable = variables.Variable(array_ops.ones([2, 3]))
@@ -2410,7 +2558,7 @@
cond=cond, body=body, loop_vars=loop_vars)
cost = math_ops.reduce_sum(tensors[2])
grad = gradients_impl.gradients(cost, [variable])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(np.ones([2, 3]), sess.run(grad[0]))
@test_util.run_deprecated_v1
@@ -2433,6 +2581,7 @@
grad = gradients_impl.gradients(cost, [c0])
self.assertAllClose(0.0, sess.run(grad[0]))
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_SerialTwoLoops(self):
with self.cached_session():
i = constant_op.constant(0, name="i")
@@ -2449,8 +2598,9 @@
_, rx = control_flow_ops.while_loop(c, b, [i, rx], parallel_iterations=1)
r = gradients_impl.gradients([rx], x)
- self.assertAllClose(1024.0, r[0].eval())
+ self.assertAllClose(1024.0, r[0])
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_ParallelTwoLoops(self):
with self.cached_session():
i = constant_op.constant(0, name="i")
@@ -2468,8 +2618,9 @@
rx = math_ops.add(r1, r2)
r = gradients_impl.gradients([rx], x)
- self.assertAllClose(64.0, r[0].eval())
+ self.assertAllClose(64.0, r[0])
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_OneOutputWithControlDependencyOnSecond(self):
with self.cached_session():
i = constant_op.constant(0, name="i")
@@ -2513,6 +2664,7 @@
self._testNestedWhileGrad_Simple(use_gpu=False)
self._testNestedWhileGrad_Simple(use_gpu=True)
+ @test_util.run_v1_only("b/120545219")
def testNestedWhileGrad_SerialInner(self):
with self.cached_session():
v = constant_op.constant(1.0)
@@ -2560,6 +2712,7 @@
r = gradients_impl.gradients(r, v)[0]
self.assertAllClose(512.0, self.evaluate(r))
+ @test_util.run_v1_only("b/120545219")
def testNestedWhileGrad_ParallelIterations(self):
# Make sure the stack pushes and pops of an inner loop are executed in
# the sequential order of the iterations of its outer loop.
@@ -2580,7 +2733,7 @@
train_op = optimizer.minimize(math_ops.reduce_mean(math_ops.square(res)))
self.evaluate(variables.global_variables_initializer())
self.evaluate(train_op)
- self.assertAllClose(2.999, self.evaluate(var))
+ self.assertAllClose(2.999, var.read_value())
def _testWhileCondGrad_Simple(self, use_gpu):
with self.cached_session(use_gpu=use_gpu):
@@ -2641,13 +2794,15 @@
[i0.get_shape(), tensor_shape.TensorShape([None, 2])])
s = math_ops.reduce_sum(h)
- self.evaluate(variables.global_variables_initializer())
optimizer = gradient_descent.GradientDescentOptimizer(0.01)
op = optimizer.minimize(s)
+
+ self.evaluate(variables.global_variables_initializer())
self.evaluate(op)
self.assertAllClose([[0.98000002, 1.98000002]], self.evaluate(x))
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
+ @test_util.run_v1_only("b/120545219")
def testWhileWithRefsWithGradients_1(self):
with self.cached_session() as sess:
x = variables.VariableV1(0.)._ref() # pylint: disable=protected-access
@@ -2665,7 +2820,7 @@
grad_ys = [variables.VariableV1(73)._ref()] # pylint: disable=protected-access
grad = gradients_impl.gradients([r[1]], [x], grad_ys=grad_ys)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(r[0].dtype, dtypes.int32)
self.assertEqual(r[1].dtype, dtypes.float32_ref)
@@ -2677,6 +2832,7 @@
self.assertEqual(73, value_x_grad)
@test_util.disable_control_flow_v2("b/116282023 (IndexedSlices)")
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_IndexedSlices(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -2699,7 +2855,7 @@
self.assertAllClose(np.array([1024.0, 1024.0]), self.evaluate(r))
@test_util.disable_control_flow_v2("b/116328420 (SparseTensor)")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileGrad_SparseTensor(self):
with self.cached_session():
values = constant_op.constant([2.0, 4.0], name="values")
@@ -2723,6 +2879,7 @@
self.assertAllClose(np.array([1024.0, 1024.0]), self.evaluate(r))
@test_util.disable_control_flow_v2("b/115920078 (gradients)")
+ @test_util.run_v1_only("b/120545219")
def testCallGradInLoop(self):
with self.cached_session() as sess:
i0 = constant_op.constant(0)
@@ -2807,6 +2964,7 @@
self.assertEqual(32.0, self.evaluate(r))
@test_util.run_deprecated_v1
+ @test_util.disable_control_flow_v2("b/118712257")
def testWhileGrad_StopGradInside(self):
with self.cached_session():
x = constant_op.constant(3.0, name="x")
@@ -2827,6 +2985,7 @@
self.assertAllClose(156.0, self.evaluate(r))
@test_util.run_deprecated_v1
+ @test_util.disable_control_flow_v2("b/118712257")
def testWhileGrad_StopGradInsideNoShape(self):
with self.cached_session() as sess:
x = array_ops.placeholder(dtypes.float32)
@@ -2860,7 +3019,7 @@
result = functional_ops.scan(fn, np.array([1., 2., 3.], dtype=np.float32))
grad_theta = gradients_impl.gradients(result, theta)
- if not control_flow_ops.ENABLE_WHILE_V2:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
with self.assertRaisesRegexp(TypeError, "Second-order gradient"):
gradients_impl.gradients(grad_theta, theta)
grad_theta_stopped = array_ops.stop_gradient(grad_theta)
@@ -2903,6 +3062,7 @@
self.assertAllClose([0., 0.], self.evaluate(dy_dq))
@test_util.disable_control_flow_v2("b/113324949 (RefVariable)")
+ @test_util.run_v1_only("b/120545219")
def testWhileGradientWithNontrainablePath2(self):
q = variables.Variable([7., 8.])
@@ -2921,6 +3081,7 @@
self.assertAllClose([1., 1.], self.evaluate(dy_dq))
@test_util.disable_control_flow_v2("b/115920078 (gradients)")
+ @test_util.run_v1_only("b/120545219")
def testIssue16504(self):
c = constant_op.constant(np.arange(100), dtype=dtypes.float32)
w = variables.Variable(
@@ -2944,6 +3105,7 @@
grad, = gradients_impl.gradients(w, c)
self.assertIsNotNone(grad)
+ @test_util.run_v1_only("b/120545219")
def testStopGradMultiFlows(self):
with self.cached_session():
@@ -2967,9 +3129,10 @@
grads = linalg_ops.norm(gradients_impl.gradients(r, vars_)[0])
z = math_ops.add(r, array_ops.stop_gradient(math_ops.reduce_sum(grads)))
result = gradients_impl.gradients(z, vars_)[0]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(5.0, self.evaluate(result))
+ @test_util.run_v1_only("b/120545219")
def testOneValueCond(self):
with self.cached_session():
@@ -3003,6 +3166,7 @@
self.assertAllClose(4.0, i.eval(feed_dict={d: 1}))
self.assertAllClose(2.0 * math.sqrt(2), i.eval(feed_dict={d: 2}))
+ @test_util.run_v1_only("b/120545219")
def testCase(self):
with self.cached_session():
x = constant_op.constant(1)
@@ -3017,14 +3181,14 @@
x < y: f1,
x > z: f2
}, default=f3, exclusive=True)
- self.assertAllEqual(r1.eval(), 17)
+ self.assertAllEqual(r1, 17)
r2 = control_flow_ops.case([(y > z, f1), (y > x, f2)], default=f3)
- self.assertAllEqual(r2.eval(), 23)
+ self.assertAllEqual(r2, 23)
# Duplicate events can happen, first one is selected
r3 = control_flow_ops.case([(x < y, f1), (x < y, f2)], default=f3)
- self.assertAllEqual(r3.eval(), 17)
+ self.assertAllEqual(r3, 17)
# Duplicate events cause an error if exclusive = True
r4 = control_flow_ops.case(
@@ -3034,7 +3198,7 @@
# Check that the default is called if none of the others are
r5 = control_flow_ops.case({x > y: f1}, default=f3)
- self.assertAllEqual(r5.eval(), -1)
+ self.assertAllEqual(r5, -1)
ran_once = [False, False, False]
@@ -3053,8 +3217,9 @@
[(x < y, break_run_twice(0)), (x > y, break_run_twice(1))],
default=lambda: constant_op.constant(2))
- self.assertAllEqual(r6.eval(), 0)
+ self.assertAllEqual(r6, 0)
+ @test_util.run_v1_only("b/120545219")
def testCaseSideEffects(self):
with self.cached_session() as sess:
v0 = variables.Variable(-1)
@@ -3075,22 +3240,23 @@
r2 = control_flow_ops.case(
((x > y, a), (x > y, b)), default=c, exclusive=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(self.evaluate([v0, v1, v2]), [-1] * 3)
self.assertEqual(2, self.evaluate(r2))
self.assertAllEqual(self.evaluate([v0, v1, v2]), [-1, -1, 2])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(self.evaluate([v0, v1, v2]), [-1] * 3)
self.assertEqual(1, self.evaluate(r1))
self.assertAllEqual(self.evaluate([v0, v1, v2]), [-1, 1, -1])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(self.evaluate([v0, v1, v2]), [-1] * 3)
self.assertEqual(0, self.evaluate(r0))
self.assertAllEqual(self.evaluate([v0, v1, v2]), [0, -1, -1])
@test_util.disable_control_flow_v2("b/113324949 (ref vars)")
+ @test_util.run_v1_only("b/120545219")
def testOneOpCond(self):
with self.cached_session():
v = variables.Variable(0)
@@ -3107,7 +3273,7 @@
i = control_flow_ops.cond(p, a, b)
self.assertTrue(isinstance(i, ops.Tensor))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(0, self.evaluate(v))
@@ -3119,6 +3285,7 @@
self.assertEqual(2, i.eval(feed_dict={c.name: 0}))
self.assertEqual(2, self.evaluate(v))
+ @test_util.run_v1_only("b/120545219")
def testWithOpsDependencies(self):
with self.cached_session() as sess:
v = variables.VariableV1(0.0)
@@ -3142,6 +3309,7 @@
# Ensure that 'v' is initialized
self.assertAllClose(0.0, real_v_val)
+ @test_util.run_v1_only("b/120545219")
def testWithTensorDependencies(self):
with self.cached_session():
v = variables.VariableV1(0.0)
@@ -3168,6 +3336,7 @@
# Ensure that 'v' is initialized
self.assertAllClose(0.0, self.evaluate(v))
+ @test_util.run_v1_only("b/120545219")
def testWithIndexedSlicesDependencies(self):
with self.cached_session():
v = variables.VariableV1(
@@ -3214,6 +3383,7 @@
self.assertDeviceEqual("", with_vdef_dep.device)
self.assertEqual([b"loc:@vdef"], with_vdef_dep.op.colocation_groups())
+ @test_util.run_v1_only("b/120545219")
def testGroup(self):
with self.cached_session() as sess:
v1 = variables.VariableV1([0.0])
@@ -3233,6 +3403,7 @@
self.assertAllClose([0.0], v1_val)
self.assertAllClose([1.0], v2_val)
+ @test_util.run_v1_only("b/120545219")
def testGroupEmpty(self):
op = control_flow_ops.group()
self.assertEqual(op.type, "NoOp")
@@ -3293,7 +3464,7 @@
self.assertEqual([None, None], m.get_shape().as_list())
self.assertEqual([], index.get_shape())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRefSelect(self):
index = array_ops.placeholder(dtypes.int32)
@@ -3348,6 +3519,7 @@
with self.assertRaises(ValueError):
sess.run(tensor_list[0])
+ @test_util.run_v1_only("b/120545219")
def testWhilePyFuncBasic(self):
def func(x):
@@ -3359,8 +3531,9 @@
lambda i, v: [i + 1, script_ops.py_func(func, [v], [dtypes.float32])[0]],
[constant_op.constant(0), constant_op.constant(2.0, dtypes.float32)],
[tensor_shape.unknown_shape(), tensor_shape.unknown_shape()])
- self.assertEqual(r[1].eval(), 65536.0)
+ self.assertEqual(self.evaluate(r[1]), 65536.0)
+ @test_util.run_v1_only("b/120545219")
def testWhileFuncBasic(self):
@function.Defun(dtypes.float32)
@@ -3374,17 +3547,17 @@
[constant_op.constant(0), x],
[tensor_shape.unknown_shape(),
tensor_shape.unknown_shape()])
- self.assertEqual(r[1].eval(), 65536.0)
-
- r = gradients_impl.gradients(r, x)[0]
- self.assertEqual(r.eval(), 524288.0)
+ grad = gradients_impl.gradients(r, x)[0]
+ self.assertEqual(self.evaluate(r[1]), 65536.0)
+ self.assertEqual(self.evaluate(grad), 524288.0)
# while_v2 does not have stacks.
- if not control_flow_ops.ENABLE_WHILE_V2:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
self.assertEqual(
len([op for op in x.graph.get_operations() if op.type == "StackV2"
]), 1)
+ @test_util.run_v1_only("b/120545219")
def testQIntSwitchMerge(self):
with self.cached_session(force_gpu=test.is_gpu_available()) as sess:
constant_qint = constant_op.constant(np.array([42]), dtypes.qint8)
@@ -3393,6 +3566,7 @@
result = control_flow_ops.merge([v_f, v_t])
self.evaluate(result)
+ @test_util.run_v1_only("b/120545219")
def testQIntRefSwitchMerge(self):
with self.cached_session(use_gpu=test.is_gpu_available()) as sess:
var_qint = gen_state_ops.variable(
@@ -3406,6 +3580,7 @@
result = control_flow_ops.ref_merge([v_f, v_t])
self.evaluate(result)
+ @test_util.run_v1_only("b/120545219")
def testUInt64SwitchMerge(self):
with self.cached_session(force_gpu=test.is_gpu_available()) as sess:
constant_uint64 = constant_op.constant(np.array([42]), dtypes.uint64)
@@ -3453,6 +3628,7 @@
math_ops.less(1, 2), true_fn, lambda: constant_op.constant(0))
return cond_tensor[0]
+ @test_util.run_v1_only("b/120545219")
def testInvalidContext(self):
# Accessing a while loop tensor outside of control flow is illegal.
while_tensor = self._getWhileTensor()
@@ -3462,7 +3638,7 @@
"is in a while loop. See info log for more details."):
math_ops.add(1, while_tensor)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInvalidContextInCond(self):
# Accessing a while loop tensor in cond is illegal.
while_tensor = self._getWhileTensor()
@@ -3475,6 +3651,7 @@
math_ops.less(1, 2), lambda: math_ops.add(1, while_tensor),
lambda: constant_op.constant(0))
+ @test_util.run_v1_only("b/120545219")
def testInvalidContextInWhile(self):
# Accessing a while loop tensor in a different while loop is illegal.
while_tensor = self._getWhileTensor()
@@ -3509,6 +3686,7 @@
control_flow_ops.cond(math_ops.less(1, 2), branch_fn, branch_fn)
+ @test_util.run_v1_only("b/120545219")
def testValidWhileContext(self):
# Accessing a tensor in a nested while is OK.
def body(_):
@@ -3517,6 +3695,7 @@
control_flow_ops.while_loop(lambda i: i < 5, body, [0])
+ @test_util.run_v1_only("b/120545219")
def testValidNestedContexts(self):
# Accessing a tensor from a cond context in a while context, all inside an
# outer while context, is OK.
@@ -3531,7 +3710,7 @@
control_flow_ops.while_loop(lambda i: i < 5, body, [0])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInvalidNestedContexts(self):
# Accessing a tensor from a while context in a different while context, all
# inside a cond context, is illegal.
@@ -3550,6 +3729,7 @@
class TupleTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testTensors(self):
for v1_first in [True, False]:
with self.cached_session():
@@ -3580,7 +3760,7 @@
self.assertAllClose([30.0], self.evaluate(t2))
self.assertAllClose([1.0], self.evaluate(v1))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testIndexedSlices(self):
for v1_first in [True, False]:
with self.cached_session():
@@ -3733,7 +3913,7 @@
with session.Session() as sess, ops.device(default_device):
# Get the initial id i, input x, and kernel.
i, x, kernel = self._getInitVariables()
- self.evaluate(variables.global_variables_initializer())
+ variables.global_variables_initializer().run()
if static_unroll:
for _ in xrange(steps):
@@ -3832,6 +4012,7 @@
isum(tensor, maximum_iterations=3).numpy(),
[1 + 3, 2 + 3, 3 + 3, 4 + 3, 5 + 3])
+ @test_util.run_v1_only("b/120545219")
def testWhileWithMaximumIterationsAndSingleArgument(self):
with context.eager_mode():
tensor = constant_op.constant(0)
@@ -3854,6 +4035,7 @@
self.assertAllEqual(t1.numpy(), tup1.numpy())
self.assertAllEqual(t2.numpy(), tup2.numpy())
+ @test_util.run_v1_only("b/120545219")
def testCase(self):
with context.eager_mode():
x = constant_op.constant(1)
diff --git a/tensorflow/python/kernel_tests/control_flow_util_test.py b/tensorflow/python/kernel_tests/control_flow_util_test.py
index 762c445..573f4b0 100644
--- a/tensorflow/python/kernel_tests/control_flow_util_test.py
+++ b/tensorflow/python/kernel_tests/control_flow_util_test.py
@@ -22,6 +22,7 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_ops
+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 control_flow_util
@@ -32,6 +33,7 @@
class ControlFlowUtilTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testIsSwitch(self):
switch_false, _ = control_flow_ops.switch(1, True)
switch = switch_false.op
@@ -44,6 +46,7 @@
self.assertFalse(control_flow_util.IsSwitch(test_ops.int_output().op))
+ @test_util.run_v1_only("b/120545219")
def testIsLoopEnter(self):
enter = gen_control_flow_ops.enter(1, frame_name="name").op
self.assertTrue(control_flow_util.IsLoopEnter(enter))
@@ -61,6 +64,7 @@
self.assertFalse(control_flow_util.IsLoopEnter(test_ops.int_output().op))
+ @test_util.run_v1_only("b/120545219")
def testIsLoopExit(self):
exit_op = control_flow_ops.exit(1).op
self.assertTrue(control_flow_util.IsLoopExit(exit_op))
diff --git a/tensorflow/python/kernel_tests/control_flow_util_v2_test.py b/tensorflow/python/kernel_tests/control_flow_util_v2_test.py
index d0374a7..08d3214 100644
--- a/tensorflow/python/kernel_tests/control_flow_util_v2_test.py
+++ b/tensorflow/python/kernel_tests/control_flow_util_v2_test.py
@@ -23,6 +23,7 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import control_flow_util_v2
from tensorflow.python.platform import test
@@ -30,14 +31,11 @@
class ControlFlowUtilV2Test(test.TestCase):
def setUp(self):
- self._enable_cond_v2_old = control_flow_ops.ENABLE_COND_V2
- self._enable_while_v2_old = control_flow_ops.ENABLE_WHILE_V2
- control_flow_ops.ENABLE_COND_V2 = True
- control_flow_ops.ENABLE_WHILE_V2 = True
+ self._enable_control_flow_v2_old = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
def tearDown(self):
- control_flow_ops.ENABLE_COND_V2 = self._enable_cond_v2_old
- control_flow_ops.ENABLE_WHILE_V2 = self._enable_while_v2_old
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = self._enable_control_flow_v2_old
def _create_control_flow(self, expect_in_defun):
"""Helper method for testInDefun."""
diff --git a/tensorflow/python/kernel_tests/ctc_loss_op_test.py b/tensorflow/python/kernel_tests/ctc_loss_op_test.py
index e6b5835..352dede 100644
--- a/tensorflow/python/kernel_tests/ctc_loss_op_test.py
+++ b/tensorflow/python/kernel_tests/ctc_loss_op_test.py
@@ -106,7 +106,7 @@
with self.assertRaisesOpError(expected_err_re):
self.evaluate([loss, grad])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBasic(self):
"""Test two batch entries."""
# Input and ground truth from Alex Graves' implementation.
@@ -242,7 +242,6 @@
self._testCTCLoss(inputs, seq_lens, labels, loss_truth, grad_truth)
- @test_util.run_deprecated_v1
def test_time_major(self):
"""Testing time_major param.
@@ -272,7 +271,7 @@
(tf_loss, tf_loss_transposed) = self.evaluate([loss, loss_transposed])
self.assertAllEqual(tf_loss, tf_loss_transposed)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInvalidSecondGradient(self):
inputs = np.random.randn(2, 2, 3).astype(np.float32)
inputs_t = constant_op.constant(inputs)
@@ -289,7 +288,7 @@
"explicitly disabled"):
_ = gradients_impl._hessian_vector_product(loss, [inputs_t], v)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testEmptyBatch(self):
inputs = constant_op.constant([], dtype=dtypes.float32, shape=(1, 0, 2))
sequence_lengths = constant_op.constant([], dtype=dtypes.int32)
@@ -306,7 +305,7 @@
class CTCLossTestV2(test.TestCase):
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCtcLossV2(self):
random_seed.set_random_seed(5)
@@ -351,7 +350,7 @@
logit_length=logit_length,
blank_index=0))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCtcLossDenseIsSameAsCtcLoss(self):
with ops.device("/GPU:0" if test.is_gpu_available() else "/CPU:0"):
random_seed.set_random_seed(5)
@@ -405,7 +404,7 @@
rtol=2e-06,
atol=2e-06)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCtcLossDenseUniqueFastPathIsSameAsCtcLoss(self):
random_seed.set_random_seed(5)
@@ -459,7 +458,7 @@
rtol=2e-06,
atol=2e-06)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCtcLossDenseWithBlankIndexIsSameAsCtcLoss(self):
random_seed.set_random_seed(5)
@@ -516,7 +515,7 @@
rtol=2e-06,
atol=2e-06)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCtcLossDenseWithNegativeBlankIndexIsSameAsCtcLoss(self):
with ops.device("/GPU:0" if test.is_gpu_available() else "/CPU:0"):
random_seed.set_random_seed(5)
@@ -565,7 +564,6 @@
rtol=2e-06,
atol=2e-06)
- @test_util.run_deprecated_v1
def testCollapseRepeated(self):
collapsed, new_seq_lengths = ctc_ops.collapse_repeated(
labels=[[1, 3, 3, 3, 0],
@@ -579,7 +577,6 @@
[1, 4, 0, 0],
[4, 2, 9, 4]])
- @test_util.run_deprecated_v1
def testCollapseRepeatedPreservesDtypes(self):
collapsed, new_seq_lengths = ctc_ops.collapse_repeated(
labels=constant_op.constant(
@@ -597,7 +594,6 @@
[1, 4, 0, 0],
[4, 2, 9, 4]])
- @test_util.run_deprecated_v1
def testCollapseRepeatedExtraPadding(self):
collapsed, new_seq_lengths = ctc_ops.collapse_repeated(
labels=[[1, 3, 3, 3, 0, 0, 0],
@@ -611,7 +607,6 @@
[1, 4, 0, 0],
[4, 2, 9, 4]])
- @test_util.run_deprecated_v1
def testCollapseRepeatedFrontRepeats(self):
collapsed, new_seq_lengths = ctc_ops.collapse_repeated(
labels=[[1, 1, 1, 2, 2],
@@ -625,7 +620,6 @@
[1, 2],
[1, 0]])
- @test_util.run_deprecated_v1
def testCollapseRepeatedAllLabelsTheSame(self):
collapsed, new_seq_lengths = ctc_ops.collapse_repeated(
labels=[[1, 1, 1, 1, 1],
@@ -658,7 +652,6 @@
self.assertAllEqual(padded_dense, new_dense)
- @test_util.run_deprecated_v1
def testUnique(self):
labels = [
[3, 4, 4, 3],
@@ -674,7 +667,6 @@
[0, 0, 0, 1],
], idx)
- @test_util.run_deprecated_v1
def testSumStates(self):
idx = [
[0, 1, 0, 1],
@@ -694,7 +686,6 @@
[1.8, 0.8, 0.0, 0.0]]
], sum_of_states)
- @test_util.run_deprecated_v1
def testStateToOlabel(self):
labels = [
[3, 4, 3, 4],
@@ -733,7 +724,6 @@
[22.0 + 23.0 + 24.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],
])
- @test_util.run_deprecated_v1
def testStateToOlabelUnique(self):
labels = [
[3, 4, 3, 4],
@@ -811,7 +801,7 @@
x = random_ops.random_uniform([])
fn = lambda accum, elem: accum + x * elem
out = ctc_ops._scan(fn, constant_op.constant([0.0, 1.0, 2.0]), 23.0)
- self.assertAllEqual(*sess.run([
+ self.assertAllClose(*sess.run([
[23.0 + x * 0.0, 23.0 + x * 1.0, 23.0 + x * 3.0], out
]))
diff --git a/tensorflow/python/kernel_tests/cwise_ops_test.py b/tensorflow/python/kernel_tests/cwise_ops_test.py
index 9bb7d8b..70f19f9 100644
--- a/tensorflow/python/kernel_tests/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/cwise_ops_test.py
@@ -887,7 +887,7 @@
tf_angle = math_ops.angle(inx)
tf_angle_val = self.evaluate(tf_angle)
- self.assertAllEqual(np_angle, tf_angle_val)
+ self.assertAllClose(np_angle, tf_angle_val)
self.assertShapeEqual(np_angle, tf_angle)
def testAngle64(self):
@@ -895,18 +895,14 @@
imag = (np.arange(-3, 3) / 5.).reshape([1, 3, 2]).astype(np.float32)
cplx = real + 1j * imag
self._compareAngle(cplx, use_gpu=False)
- # TODO: Enable GPU tests for angle op after resolving
- # build failures on GPU (See #10643 for context).
- # self._compareAngle(cplx, use_gpu=True)
+ self._compareAngle(cplx, use_gpu=True)
def testAngle(self):
real = (np.arange(-3, 3) / 4.).reshape([1, 3, 2]).astype(np.float64)
imag = (np.arange(-3, 3) / 5.).reshape([1, 3, 2]).astype(np.float64)
cplx = real + 1j * imag
self._compareAngle(cplx, use_gpu=False)
- # TODO: Enable GPU tests for angle op after resolving
- # build failures on GPU (See #10643 for context).
- # self._compareAngle(cplx, use_gpu=True)
+ self._compareAngle(cplx, use_gpu=True)
@test_util.run_deprecated_v1
def testRealReal(self):
diff --git a/tensorflow/python/kernel_tests/denormal_test.py b/tensorflow/python/kernel_tests/denormal_test.py
index 80a3033..d824e95 100644
--- a/tensorflow/python/kernel_tests/denormal_test.py
+++ b/tensorflow/python/kernel_tests/denormal_test.py
@@ -36,8 +36,9 @@
self.assertEqual(tiny, tiny / 16 * 16)
def _flushDenormalsTest(self, use_gpu, dtypes):
- if platform.machine() == "ppc64le" or platform.machine() == "s390x":
- # Disabled denormal_test on power/s390x platform
+ if platform.machine() == "ppc64le" or platform.machine(
+ ) == "s390x" or platform.machine() == "aarch64":
+ # Disabled denormal_test on power/s390x/aarch64 platform
# Check relevant discussion - https://github.com/tensorflow/tensorflow/issues/11902
return
with self.cached_session(use_gpu=use_gpu):
diff --git a/tensorflow/python/kernel_tests/dense_update_ops_no_tsan_test.py b/tensorflow/python/kernel_tests/dense_update_ops_no_tsan_test.py
index 4f74e1e..a778bf2 100644
--- a/tensorflow/python/kernel_tests/dense_update_ops_no_tsan_test.py
+++ b/tensorflow/python/kernel_tests/dense_update_ops_no_tsan_test.py
@@ -20,7 +20,6 @@
import numpy as np
-from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
@@ -33,7 +32,6 @@
# NOTE(mrry): We exclude thess tests from the TSAN TAP target, because they
# contain benign and deliberate data races when multiple threads update
# the same parameters without a lock.
- @test_util.run_deprecated_v1
def testParallelUpdateWithoutLocking(self):
with self.cached_session() as sess:
ones_t = array_ops.fill([1024, 1024], 1.0)
@@ -42,7 +40,7 @@
state_ops.assign_add(
p, ones_t, use_locking=False) for _ in range(20)
]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
def run_add(add_op):
self.evaluate(add_op)
@@ -61,7 +59,6 @@
self.assertTrue((vals >= ones).all())
self.assertTrue((vals <= ones * 20).all())
- @test_util.run_deprecated_v1
def testParallelAssignWithoutLocking(self):
with self.cached_session() as sess:
ones_t = array_ops.fill([1024, 1024], float(1))
@@ -70,7 +67,7 @@
state_ops.assign(p, math_ops.multiply(ones_t, float(i)), False)
for i in range(1, 21)
]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
def run_assign(assign_op):
self.evaluate(assign_op)
@@ -94,7 +91,6 @@
# contain non-benign but known data races between the variable assignment and
# returning the output tensors. This issue will be resolved with the new
# resource variables.
- @test_util.run_deprecated_v1
def testParallelUpdateWithLocking(self):
with self.cached_session() as sess:
zeros_t = array_ops.fill([1024, 1024], 0.0)
@@ -104,7 +100,7 @@
state_ops.assign_add(
p, ones_t, use_locking=True) for _ in range(20)
]
- p.initializer.run()
+ self.evaluate(p.initializer)
def run_add(add_op):
self.evaluate(add_op)
@@ -122,7 +118,6 @@
ones = np.ones((1024, 1024)).astype(np.float32)
self.assertAllEqual(vals, ones * 20)
- @test_util.run_deprecated_v1
def testParallelAssignWithLocking(self):
with self.cached_session() as sess:
zeros_t = array_ops.fill([1024, 1024], 0.0)
@@ -133,7 +128,7 @@
p, math_ops.multiply(ones_t, float(i)), use_locking=True)
for i in range(1, 21)
]
- p.initializer.run()
+ self.evaluate(p.initializer)
def run_assign(assign_op):
self.evaluate(assign_op)
diff --git a/tensorflow/python/kernel_tests/dense_update_ops_test.py b/tensorflow/python/kernel_tests/dense_update_ops_test.py
index 309da88..545de87 100644
--- a/tensorflow/python/kernel_tests/dense_update_ops_test.py
+++ b/tensorflow/python/kernel_tests/dense_update_ops_test.py
@@ -86,7 +86,7 @@
def testBasic(self):
self._testTypes(np.arange(0, 20).reshape([4, 5]))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAssignNonStrictShapeChecking(self):
with self.cached_session():
data = array_ops.fill([1024, 1024], 0)
@@ -101,7 +101,7 @@
a2.op.run()
self.assertAllEqual(p.eval(), self.evaluate(data2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitRequiredAssignAdd(self):
with self.cached_session():
p = variables.VariableV1(array_ops.fill([1024, 1024], 1), dtypes.int32)
@@ -109,7 +109,7 @@
with self.assertRaisesOpError("use uninitialized"):
a.op.run()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitRequiredAssignSub(self):
with self.cached_session():
p = variables.VariableV1(array_ops.fill([1024, 1024], 1), dtypes.int32)
diff --git a/tensorflow/python/kernel_tests/depthwise_conv_op_test.py b/tensorflow/python/kernel_tests/depthwise_conv_op_test.py
index f6d834c..5b1a47f 100644
--- a/tensorflow/python/kernel_tests/depthwise_conv_op_test.py
+++ b/tensorflow/python/kernel_tests/depthwise_conv_op_test.py
@@ -24,6 +24,7 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradient_checker
from tensorflow.python.ops import nn_impl
@@ -185,6 +186,7 @@
self.assertShapeEqual(native_result, conv_native)
self.assertShapeEqual(native_result, conv_interface)
+ @test_util.run_v1_only("b/120545219")
def testDepthwiseConv2D(self):
for index, (input_size, filter_size, _, stride,
padding) in enumerate(ConfigsToTest()):
@@ -428,6 +430,7 @@
use_gpu, grouped_conv, err)
self.assertLess(err, tolerance)
+ @test_util.run_v1_only("b/120545219")
def testDepthwiseConv2DInputGrad(self):
for index, (input_size, filter_size, output_size, stride,
padding) in enumerate(CheckGradConfigsToTest()):
@@ -477,6 +480,7 @@
use_gpu=True,
data_format="NCHW")
+ @test_util.run_v1_only("b/120545219")
def testDepthwiseConv2DFilterGrad(self):
for index, (input_size, filter_size, output_size, stride,
padding) in enumerate(CheckGradConfigsToTest()):
diff --git a/tensorflow/python/kernel_tests/determinant_op_test.py b/tensorflow/python/kernel_tests/determinant_op_test.py
index d6ef9e7..dbfda38 100644
--- a/tensorflow/python/kernel_tests/determinant_op_test.py
+++ b/tensorflow/python/kernel_tests/determinant_op_test.py
@@ -133,6 +133,7 @@
huge_matrix = np.array([[max_double, 0.0], [0.0, max_double]])
self._compareDeterminant(huge_matrix)
+ @test_util.run_v1_only("b/120545219")
def testNonSquareMatrix(self):
# When the determinant of a non-square matrix is attempted we should return
# an error
@@ -140,6 +141,7 @@
linalg_ops.matrix_determinant(
np.array([[1., 2., 3.], [3., 5., 4.]]).astype(np.float32))
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
# The input to the determinant should be a 2-dimensional tensor.
tensor1 = constant_op.constant([1., 2.])
@@ -150,6 +152,7 @@
self._compareDeterminant(np.empty([0, 2, 2]))
self._compareDeterminant(np.empty([2, 0, 0]))
+ @test_util.run_v1_only("b/120545219")
def testConcurrentExecutesWithoutError(self):
with self.session(use_gpu=True) as sess:
matrix1 = random_ops.random_normal([5, 5], seed=42)
diff --git a/tensorflow/python/kernel_tests/distributions/multinomial_test.py b/tensorflow/python/kernel_tests/distributions/multinomial_test.py
index b3f3416..187ddd4 100644
--- a/tensorflow/python/kernel_tests/distributions/multinomial_test.py
+++ b/tensorflow/python/kernel_tests/distributions/multinomial_test.py
@@ -22,6 +22,7 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.distributions import multinomial
@@ -33,6 +34,7 @@
def setUp(self):
self._rng = np.random.RandomState(42)
+ @test_util.run_v1_only("b/120545219")
def testSimpleShapes(self):
with self.cached_session():
p = [.1, .3, .6]
@@ -42,6 +44,7 @@
self.assertEqual(tensor_shape.TensorShape([3]), dist.event_shape)
self.assertEqual(tensor_shape.TensorShape([]), dist.batch_shape)
+ @test_util.run_v1_only("b/120545219")
def testComplexShapes(self):
with self.cached_session():
p = 0.5 * np.ones([3, 2, 2], dtype=np.float32)
@@ -52,6 +55,7 @@
self.assertEqual(tensor_shape.TensorShape([2]), dist.event_shape)
self.assertEqual(tensor_shape.TensorShape([3, 2]), dist.batch_shape)
+ @test_util.run_v1_only("b/120545219")
def testN(self):
p = [[0.1, 0.2, 0.7], [0.2, 0.3, 0.5]]
n = [[3.], [4]]
@@ -60,6 +64,7 @@
self.assertEqual((2, 1), dist.total_count.get_shape())
self.assertAllClose(n, dist.total_count.eval())
+ @test_util.run_v1_only("b/120545219")
def testP(self):
p = [[0.1, 0.2, 0.7]]
with self.cached_session():
@@ -68,6 +73,7 @@
self.assertEqual((1, 3), dist.logits.get_shape())
self.assertAllClose(p, dist.probs.eval())
+ @test_util.run_v1_only("b/120545219")
def testLogits(self):
p = np.array([[0.1, 0.2, 0.7]], dtype=np.float32)
logits = np.log(p) - 50.
@@ -78,6 +84,7 @@
self.assertAllClose(p, multinom.probs.eval())
self.assertAllClose(logits, multinom.logits.eval())
+ @test_util.run_v1_only("b/120545219")
def testPmfUnderflow(self):
logits = np.array([[-200, 0]], dtype=np.float32)
with self.cached_session():
@@ -85,6 +92,7 @@
lp = dist.log_prob([1., 0.]).eval()[0]
self.assertAllClose(-200, lp, atol=0, rtol=1e-6)
+ @test_util.run_v1_only("b/120545219")
def testPmfandCountsAgree(self):
p = [[0.1, 0.2, 0.7]]
n = [[5.]]
@@ -97,6 +105,7 @@
with self.assertRaisesOpError("counts must sum to `self.total_count`"):
dist.prob([3., 3, 0]).eval()
+ @test_util.run_v1_only("b/120545219")
def testPmfNonIntegerCounts(self):
p = [[0.1, 0.2, 0.7]]
n = [[5.]]
@@ -157,6 +166,7 @@
self.assertAllClose([0.1, 0.9], self.evaluate(pmf))
self.assertEqual((2), pmf.get_shape())
+ @test_util.run_v1_only("b/120545219")
def testPmfCountsStretchedInBroadcastWhenSameRank(self):
with self.cached_session():
p = [[0.1, 0.9], [0.7, 0.3]]
@@ -165,6 +175,7 @@
self.assertAllClose(pmf.eval(), [0.1, 0.7])
self.assertEqual((2), pmf.get_shape())
+ @test_util.run_v1_only("b/120545219")
def testPmfCountsStretchedInBroadcastWhenLowerRank(self):
with self.cached_session():
p = [[0.1, 0.9], [0.7, 0.3]]
@@ -194,6 +205,7 @@
self.evaluate(pmf)
self.assertEqual((4, 3), pmf.get_shape())
+ @test_util.run_v1_only("b/120545219")
def testMultinomialMean(self):
with self.cached_session():
n = 5.
@@ -203,6 +215,7 @@
self.assertEqual((3,), dist.mean().get_shape())
self.assertAllClose(expected_means, dist.mean().eval())
+ @test_util.run_v1_only("b/120545219")
def testMultinomialCovariance(self):
with self.cached_session():
n = 5.
@@ -214,6 +227,7 @@
self.assertEqual((3, 3), dist.covariance().get_shape())
self.assertAllClose(expected_covariances, dist.covariance().eval())
+ @test_util.run_v1_only("b/120545219")
def testMultinomialCovarianceBatch(self):
with self.cached_session():
# Shape [2]
@@ -246,6 +260,7 @@
self.assertEqual((3, 5, 4, 4), covariance.get_shape())
self.assertEqual((6, 3, 3, 3), covariance2.get_shape())
+ @test_util.run_v1_only("b/120545219")
def testCovarianceFromSampling(self):
# We will test mean, cov, var, stddev on a DirichletMultinomial constructed
# via broadcast between alpha, n.
@@ -288,6 +303,7 @@
self.assertAllClose(sample_var_, analytic_var, atol=0.01, rtol=0.01)
self.assertAllClose(sample_stddev_, analytic_stddev, atol=0.01, rtol=0.01)
+ @test_util.run_v1_only("b/120545219")
def testSampleUnbiasedNonScalarBatch(self):
with self.cached_session() as sess:
dist = multinomial.Multinomial(
@@ -317,6 +333,7 @@
self.assertAllClose(
actual_covariance_, sample_covariance_, atol=0., rtol=0.20)
+ @test_util.run_v1_only("b/120545219")
def testSampleUnbiasedScalarBatch(self):
with self.cached_session() as sess:
dist = multinomial.Multinomial(
diff --git a/tensorflow/python/kernel_tests/embedding_ops_test.py b/tensorflow/python/kernel_tests/embedding_ops_test.py
index 6019245..3ea2071 100644
--- a/tensorflow/python/kernel_tests/embedding_ops_test.py
+++ b/tensorflow/python/kernel_tests/embedding_ops_test.py
@@ -851,8 +851,9 @@
embedding_weights = self._random_weights()
sparse_ids, sparse_weights = self._ids_and_weights_2d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, sparse_weights).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, sparse_weights).eval())
self.assertAllClose(
embedding_lookup_result,
@@ -865,8 +866,10 @@
embedding_weights = self._random_weights()
sparse_ids, sparse_weights = self._ids_and_weights_2d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, sparse_weights, default_id=3).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, sparse_weights,
+ default_id=3).eval())
self.assertAllClose(
embedding_lookup_result,
@@ -880,8 +883,9 @@
embedding_weights = self._random_weights()
sparse_ids, _ = self._ids_and_weights_2d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, None).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, None).eval())
self.assertAllClose(
embedding_lookup_result,
@@ -895,8 +899,9 @@
embedding_weights = self._random_weights(num_shards=3)
sparse_ids, _ = self._ids_and_weights_2d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, None).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, None).eval())
embedding_weights = list(itertools.chain(*embedding_weights))
self.assertAllClose(embedding_lookup_result,
@@ -926,8 +931,9 @@
embedding_weights = self._random_weights()
sparse_ids, sparse_weights = self._ids_and_weights_3d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, sparse_weights).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, sparse_weights).eval())
self.assertAllClose(embedding_lookup_result, [[
(1.0 * embedding_weights[0][0] + 2.0 * embedding_weights[0][1]) / 3.0,
@@ -940,8 +946,10 @@
embedding_weights = self._random_weights()
sparse_ids, sparse_weights = self._ids_and_weights_3d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, sparse_weights, default_id=3).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, sparse_weights,
+ default_id=3).eval())
self.assertAllClose(
embedding_lookup_result,
@@ -957,8 +965,9 @@
embedding_weights = self._random_weights()
sparse_ids, _ = self._ids_and_weights_3d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, None).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, None).eval())
self.assertAllClose(embedding_lookup_result, [[(
embedding_weights[0][0] + embedding_weights[0][1]) / 2.0, [0] * 4, [
@@ -974,8 +983,9 @@
embedding_weights = self._random_weights(num_shards=3)
sparse_ids, _ = self._ids_and_weights_3d()
- embedding_lookup_result = (embedding_ops.safe_embedding_lookup_sparse(
- embedding_weights, sparse_ids, None).eval())
+ embedding_lookup_result = (
+ embedding_ops.safe_embedding_lookup_sparse_v2(
+ embedding_weights, sparse_ids, None).eval())
embedding_weights = list(itertools.chain(*embedding_weights))
self.assertAllClose(embedding_lookup_result, [[
diff --git a/tensorflow/python/kernel_tests/fifo_queue_test.py b/tensorflow/python/kernel_tests/fifo_queue_test.py
index 9655351..0579ddd 100644
--- a/tensorflow/python/kernel_tests/fifo_queue_test.py
+++ b/tensorflow/python/kernel_tests/fifo_queue_test.py
@@ -39,6 +39,7 @@
from tensorflow.python.util import compat
+@test_util.run_v1_only("b/120545219")
class FIFOQueueTest(test.TestCase):
def testConstructor(self):
@@ -1423,6 +1424,7 @@
session.run([a, c])
+@test_util.run_v1_only("b/120545219")
class FIFOQueueDictTest(test.TestCase):
def testConstructor(self):
@@ -1583,6 +1585,7 @@
self.assertTrue([compat.as_bytes("dd"), compat.as_bytes("ee")], list(s))
+@test_util.run_v1_only("b/120545219")
class FIFOQueueWithTimeoutTest(test.TestCase):
def testDequeueWithTimeout(self):
@@ -1617,6 +1620,7 @@
self.assertEqual(37, self.evaluate(dequeued_t))
+@test_util.run_v1_only("b/120545219")
class QueueContainerTest(test.TestCase):
def testContainer(self):
diff --git a/tensorflow/python/kernel_tests/functional_ops_test.py b/tensorflow/python/kernel_tests/functional_ops_test.py
index c489623..0d6a3cb 100644
--- a/tensorflow/python/kernel_tests/functional_ops_test.py
+++ b/tensorflow/python/kernel_tests/functional_ops_test.py
@@ -466,7 +466,7 @@
loss = l0 + array_ops.stop_gradient(l1)
grad = gradients_impl.gradients(ys=[loss], xs=[a, b])
with self.test_session(use_gpu=True) as sess:
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.evaluate(grad)
@test_util.run_in_graph_and_eager_modes
@@ -494,7 +494,7 @@
@test_util.disable_control_flow_v2("b/119323354")
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMapEmptyScalar(self):
map_return = functional_ops.map_fn(lambda x: 1, constant_op.constant([]))
self.assertAllEqual([0], map_return.get_shape().dims)
@@ -503,7 +503,7 @@
# TODO(akshayka): this test fails in eager: the iterable is of length 0 so
# so the body of the while loop never executes
@test_util.disable_control_flow_v2("b/119323354")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMapEmptyTensor(self):
with self.cached_session():
map_return = functional_ops.map_fn(lambda x: array_ops.zeros([3, 2]),
@@ -797,7 +797,7 @@
self.assertAllEqual(Run(100., False), 5050.)
self.assertAllEqual(Run(100., True), 5050.)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWhileError(self):
for use_gpu in (True, False):
with ops.Graph().as_default() as g:
@@ -1027,7 +1027,7 @@
def testForMLPWhile(self):
self._testForMLP(True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testForError(self):
@function.Defun(dtypes.int32, dtypes.float32)
@@ -1233,7 +1233,7 @@
self.assertAllEqual(expected, result)
# Use an invalid executor name to test the plumbing of the executor_type attr.
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testExecutorTypeAttrExecutorNotFound(self):
@function.Defun(dtypes.int32)
def AddFive(x):
diff --git a/tensorflow/python/kernel_tests/identity_op_py_test.py b/tensorflow/python/kernel_tests/identity_op_py_test.py
index 1a6794e..40ec9db 100644
--- a/tensorflow/python/kernel_tests/identity_op_py_test.py
+++ b/tensorflow/python/kernel_tests/identity_op_py_test.py
@@ -62,7 +62,7 @@
self.assertEquals(shape,
array_ops.identity(np.array(array_2x3)).get_shape())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRefIdentityShape(self):
with self.cached_session():
shape = [2, 3]
diff --git a/tensorflow/python/kernel_tests/linalg_grad_test.py b/tensorflow/python/kernel_tests/linalg_grad_test.py
index 28e1d7e..ff84221 100644
--- a/tensorflow/python/kernel_tests/linalg_grad_test.py
+++ b/tensorflow/python/kernel_tests/linalg_grad_test.py
@@ -61,6 +61,7 @@
def _GetMatrixUnaryFunctorGradientTest(functor_, dtype_, shape_, **kwargs_):
+ @test_util.run_v1_only('b/120545219')
def Test(self):
with self.session(use_gpu=True):
np.random.seed(1)
@@ -103,6 +104,7 @@
float32_tol_fudge=1.0,
**kwargs_):
+ @test_util.run_v1_only('b/120545219')
def Test(self):
# TODO(rmlarsen): Debug illegal address bug on CUDA and re-enable
# GPU test for matrix_solve.
diff --git a/tensorflow/python/kernel_tests/losses_test.py b/tensorflow/python/kernel_tests/losses_test.py
index abff61f..4584a27 100644
--- a/tensorflow/python/kernel_tests/losses_test.py
+++ b/tensorflow/python/kernel_tests/losses_test.py
@@ -51,26 +51,26 @@
losses.absolute_difference(
self._predictions, self._predictions, weights=None)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAllCorrectNoLossWeight(self):
loss = losses.absolute_difference(self._predictions, self._predictions)
with self.cached_session():
self.assertAlmostEqual(0.0, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLoss(self):
loss = losses.absolute_difference(self._labels, self._predictions)
with self.cached_session():
self.assertAlmostEqual(5.5, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithPythonScalarWeight(self):
weights = 2.3
loss = losses.absolute_difference(self._labels, self._predictions, weights)
with self.cached_session():
self.assertAlmostEqual(5.5 * weights, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeight(self):
weights = 2.3
loss = losses.absolute_difference(self._labels, self._predictions,
@@ -148,7 +148,7 @@
self.assertEquals(loss.op.name, 'softmax_cross_entropy_loss/value')
self.assertAlmostEqual(loss.eval(), 10.0, 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithPythonScalarWeight(self):
logits = constant_op.constant([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0],
[0.0, 0.0, 10.0]])
@@ -158,7 +158,7 @@
loss = losses.softmax_cross_entropy(labels, logits, weights)
self.assertAlmostEqual(weights * 10.0, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeight(self):
logits = constant_op.constant([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0],
[0.0, 0.0, 10.0]])
@@ -311,7 +311,7 @@
self.assertEquals(loss.op.name, 'sparse_softmax_cross_entropy_loss/value')
self.assertAlmostEqual(loss.eval(), 10.0, 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithPythonScalarWeight(self):
logits = constant_op.constant([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0],
[0.0, 0.0, 10.0]])
@@ -321,7 +321,7 @@
loss = losses.sparse_softmax_cross_entropy(labels, logits, weights)
self.assertAlmostEqual(weights * 10.0, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeight(self):
logits = constant_op.constant([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0],
[0.0, 0.0, 10.0]])
@@ -677,13 +677,13 @@
with self.assertRaises(ValueError):
losses.log_loss(self._labels, self._labels, weights=None)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAllCorrectNoLossWeight(self):
loss = losses.log_loss(self._labels, self._labels)
with self.cached_session():
self.assertAlmostEqual(0.0, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAllCorrectNoLossWeightWithPlaceholder(self):
tf_predictions = array_ops.placeholder(
dtypes.float32, shape=self._np_labels.shape)
@@ -692,14 +692,14 @@
self.assertAlmostEqual(
0.0, loss.eval(feed_dict={tf_predictions: self._np_labels}), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLoss(self):
loss = losses.log_loss(self._labels, self._predictions)
with self.cached_session():
self.assertAlmostEqual(-np.sum(self._expected_losses) / 6.0,
self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithPythonScalarWeight(self):
weights = 2.3
loss = losses.log_loss(self._labels, self._predictions, weights)
@@ -707,7 +707,7 @@
self.assertAlmostEqual(weights * -np.sum(self._expected_losses) / 6.0,
self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeight(self):
weights = 2.3
loss = losses.log_loss(self._labels, self._predictions,
@@ -716,7 +716,7 @@
self.assertAlmostEqual(weights * -np.sum(self._expected_losses) / 6.0,
self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeightAndPlaceholder(self):
tf_predictions = array_ops.placeholder(
dtypes.float32, shape=self._np_predictions.shape)
@@ -728,7 +728,7 @@
self.assertAlmostEqual(weights * -np.sum(self._expected_losses) / 6.0,
loss, 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeightAndPlaceholderWithRankOnly(self):
tf_predictions = array_ops.placeholder(dtypes.float32, shape=[None, None])
weights = 2.3
@@ -788,7 +788,7 @@
self.assertAlmostEqual(-np.sum(expected_losses) / 5.0,
self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithMeasurementSpecificWeightsWithPlaceholder(self):
weights = np.array([3, 6, 5, 0, 4, 2]).reshape((2, 3))
expected_losses = np.multiply(self._expected_losses, weights)
@@ -816,7 +816,7 @@
with self.cached_session():
self.assertAlmostEqual(-np.sum(expected_losses), self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithSampleSpecificWeightsMostZeroWithPlaceholder(self):
weights = np.array([0, 0, 0, 0, 0, 2]).reshape((2, 3))
expected_losses = np.multiply(self._expected_losses, weights)
@@ -955,26 +955,26 @@
losses.mean_squared_error(predictions=constant_op.constant(0),
labels=constant_op.constant(0)).eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAllCorrectNoLossWeight(self):
loss = losses.mean_squared_error(self._predictions, self._predictions)
with self.cached_session():
self.assertAlmostEqual(0.0, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLoss(self):
loss = losses.mean_squared_error(self._labels, self._predictions)
with self.cached_session():
self.assertAlmostEqual(49.5, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithPythonScalarWeight(self):
weights = 2.3
loss = losses.mean_squared_error(self._labels, self._predictions, weights)
with self.cached_session():
self.assertAlmostEqual(49.5 * weights, self.evaluate(loss), 3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeight(self):
weights = 2.3
loss = losses.mean_squared_error(self._labels, self._predictions,
@@ -1068,12 +1068,12 @@
self.assertAlmostEqual(
expected_loss, dynamic_inputs_op.eval(feed_dict=feed_dict), places=3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAllCorrectNoLossWeight(self):
self._test_valid_weights(
self._labels, self._labels, expected_loss=0.0)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLoss(self):
self._test_valid_weights(
self._labels, self._predictions,
@@ -1104,7 +1104,7 @@
np_grad = self.evaluate(grad)
self.assertFalse(np.isnan(np_grad).any())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithPythonScalarWeight(self):
weight = 2.3
self._test_valid_weights(
@@ -1112,7 +1112,7 @@
expected_loss=weight * np.sum(self._expected_losses),
weights=weight)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonZeroLossWithScalarTensorWeight(self):
weights = 2.3
loss = losses.mean_pairwise_squared_error(
@@ -1179,7 +1179,7 @@
weights_placeholder: weights,
})
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInvalid3dWeighted2x0(self):
labels = np.array([
[[1, 9, 2], [12, 11, 10], [9, 8, 7]],
diff --git a/tensorflow/python/kernel_tests/lu_op_test.py b/tensorflow/python/kernel_tests/lu_op_test.py
new file mode 100644
index 0000000..06deb0e
--- /dev/null
+++ b/tensorflow/python/kernel_tests/lu_op_test.py
@@ -0,0 +1,288 @@
+# 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 tensorflow.ops.tf.Lu."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+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 functional_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import benchmark
+from tensorflow.python.platform import test
+
+
+class LuOpTest(test.TestCase):
+
+ @property
+ def float_types(self):
+ return set((np.float64, np.float32, np.complex64, np.complex128))
+
+ def _verifyLuBase(self, x, lower, upper, perm, verification,
+ output_idx_type):
+ lower_np, upper_np, perm_np, verification_np = self.evaluate(
+ [lower, upper, perm, verification])
+
+ self.assertAllClose(x, verification_np)
+ self.assertShapeEqual(x, lower)
+ self.assertShapeEqual(x, upper)
+
+ self.assertAllEqual(x.shape[:-1], perm.shape.as_list())
+
+ # Check dtypes are as expected.
+ self.assertEqual(x.dtype, lower_np.dtype)
+ self.assertEqual(x.dtype, upper_np.dtype)
+ self.assertEqual(output_idx_type.as_numpy_dtype, perm_np.dtype)
+
+ # Check that the permutation is valid.
+ if perm_np.shape[-1] > 0:
+ perm_reshaped = np.reshape(perm_np, (-1, perm_np.shape[-1]))
+ for perm_vector in perm_reshaped:
+ self.assertAllClose(np.arange(len(perm_vector)), np.sort(perm_vector))
+
+ def _verifyLu(self, x, output_idx_type=dtypes.int64):
+ # Verify that Px = LU.
+ with test_util.use_gpu():
+
+ lu, perm = linalg_ops.lu(x, output_idx_type=output_idx_type)
+
+ # Prepare the lower factor of shape num_rows x num_rows
+ lu_shape = np.array(lu.shape.as_list())
+ batch_shape = lu_shape[:-2]
+ num_rows = lu_shape[-2]
+ num_cols = lu_shape[-1]
+
+ lower = array_ops.matrix_band_part(lu, -1, 0)
+
+ if num_rows > num_cols:
+ eye = linalg_ops.eye(
+ num_rows, batch_shape=batch_shape, dtype=lower.dtype)
+ lower = array_ops.concat([lower, eye[..., num_cols:]], axis=-1)
+ elif num_rows < num_cols:
+ lower = lower[..., :num_rows]
+
+ # Fill the diagonal with ones.
+ ones_diag = array_ops.ones(
+ np.append(batch_shape, num_rows), dtype=lower.dtype)
+ lower = array_ops.matrix_set_diag(lower, ones_diag)
+
+ # Prepare the upper factor.
+ upper = array_ops.matrix_band_part(lu, 0, -1)
+
+ verification = math_ops.matmul(lower, upper)
+
+ # Permute the rows of product of the Cholesky factors.
+ if num_rows > 0:
+ # Reshape the product of the triangular factors and permutation indices
+ # to a single batch dimension. This makes it easy to apply
+ # invert_permutation and gather_nd ops.
+ perm_reshaped = array_ops.reshape(perm, [-1, num_rows])
+ verification_reshaped = array_ops.reshape(verification,
+ [-1, num_rows, num_cols])
+ # Invert the permutation in each batch.
+ inv_perm_reshaped = functional_ops.map_fn(array_ops.invert_permutation,
+ perm_reshaped)
+ batch_size = perm_reshaped.shape.as_list()[0]
+ # Prepare the batch indices with the same shape as the permutation.
+ # The corresponding batch index is paired with each of the `num_rows`
+ # permutation indices.
+ batch_indices = math_ops.cast(
+ array_ops.broadcast_to(
+ math_ops.range(batch_size)[:, None], perm_reshaped.shape),
+ dtype=output_idx_type)
+ permuted_verification_reshaped = array_ops.gather_nd(
+ verification_reshaped,
+ array_ops.stack([batch_indices, inv_perm_reshaped], axis=-1))
+
+ # Reshape the verification matrix back to the original shape.
+ verification = array_ops.reshape(permuted_verification_reshaped,
+ lu_shape)
+
+ self._verifyLuBase(x, lower, upper, perm, verification,
+ output_idx_type)
+
+ def testBasic(self):
+ data = np.array([[4., -1., 2.], [-1., 6., 0], [10., 0., 5.]])
+
+ for dtype in (np.float32, np.float64):
+ for output_idx_type in (dtypes.int32, dtypes.int64):
+ self._verifyLu(data.astype(dtype), output_idx_type=output_idx_type)
+
+ for dtype in (np.complex64, np.complex128):
+ for output_idx_type in (dtypes.int32, dtypes.int64):
+ complex_data = np.tril(1j * data, -1).astype(dtype)
+ complex_data += np.triu(-1j * data, 1).astype(dtype)
+ complex_data += data
+ self._verifyLu(complex_data, output_idx_type=output_idx_type)
+
+ def testPivoting(self):
+ with test_util.use_gpu():
+ # This matrix triggers partial pivoting because the first diagonal entry
+ # is small.
+ data = np.array([[1e-9, 1., 0.], [1., 0., 0], [0., 1., 5]])
+ self._verifyLu(data.astype(np.float32))
+
+ for dtype in (np.float32, np.float64):
+ self._verifyLu(data.astype(dtype))
+ _, p = linalg_ops.lu(data)
+ p_val = self.evaluate([p])
+ # Make sure p_val is not the identity permutation.
+ self.assertNotAllClose(np.arange(3), p_val)
+
+ for dtype in (np.complex64, np.complex128):
+ complex_data = np.tril(1j * data, -1).astype(dtype)
+ complex_data += np.triu(-1j * data, 1).astype(dtype)
+ complex_data += data
+ self._verifyLu(complex_data)
+ _, p = linalg_ops.lu(data)
+ p_val = self.evaluate([p])
+ # Make sure p_val is not the identity permutation.
+ self.assertNotAllClose(np.arange(3), p_val)
+
+ def testInvalidMatrix(self):
+ # LU factorization gives an error when the input is singular.
+ # Note: A singular matrix may return without error but it won't be a valid
+ # factorization.
+ with test_util.use_gpu():
+ for dtype in self.float_types:
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(
+ linalg_ops.lu(
+ np.array([[1., 2., 3.], [2., 4., 6.], [2., 3., 4.]],
+ dtype=dtype)))
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(
+ linalg_ops.lu(
+ np.array([[[1., 2., 3.], [2., 4., 6.], [1., 2., 3.]],
+ [[1., 2., 3.], [3., 4., 5.], [5., 6., 7.]]],
+ dtype=dtype)))
+
+ def testBatch(self):
+ simple_array = np.array([[[1., -1.], [2., 5.]]]) # shape (1, 2, 2)
+ self._verifyLu(simple_array)
+ self._verifyLu(np.vstack((simple_array, simple_array)))
+ odd_sized_array = np.array([[[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]]])
+ self._verifyLu(np.vstack((odd_sized_array, odd_sized_array)))
+
+ batch_size = 200
+
+ # Generate random matrices.
+ np.random.seed(42)
+ matrices = np.random.rand(batch_size, 5, 5)
+ self._verifyLu(matrices)
+
+ # Generate random complex valued matrices.
+ np.random.seed(52)
+ matrices = np.random.rand(batch_size, 5,
+ 5) + 1j * np.random.rand(batch_size, 5, 5)
+ self._verifyLu(matrices)
+
+ def testLargeMatrix(self):
+ # Generate random matrices.
+ n = 500
+ np.random.seed(64)
+ data = np.random.rand(n, n)
+ self._verifyLu(data)
+
+ # Generate random complex valued matrices.
+ np.random.seed(129)
+ data = np.random.rand(n, n) + 1j * np.random.rand(n, n)
+ self._verifyLu(data)
+
+ @test_util.run_v1_only("b/120545219")
+ def testEmpty(self):
+ self._verifyLu(np.empty([0, 2, 2]))
+ self._verifyLu(np.empty([2, 0, 0]))
+
+ @test_util.run_deprecated_v1
+ def testConcurrentExecutesWithoutError(self):
+ with test_util.use_gpu():
+ matrix1 = random_ops.random_normal([5, 5], seed=42)
+ matrix2 = random_ops.random_normal([5, 5], seed=42)
+ lu1, p1 = linalg_ops.lu(matrix1)
+ lu2, p2 = linalg_ops.lu(matrix2)
+ lu1_val, p1_val, lu2_val, p2_val = self.evaluate([lu1, p1, lu2, p2])
+ self.assertAllEqual(lu1_val, lu2_val)
+ self.assertAllEqual(p1_val, p2_val)
+
+
+class LuBenchmark(test.Benchmark):
+ shapes = [
+ (4, 4),
+ (10, 10),
+ (16, 16),
+ (101, 101),
+ (256, 256),
+ (1000, 1000),
+ (1024, 1024),
+ (2048, 2048),
+ (4096, 4096),
+ (513, 2, 2),
+ (513, 8, 8),
+ (513, 256, 256),
+ (4, 513, 2, 2),
+ ]
+
+ def _GenerateMatrix(self, shape):
+ batch_shape = shape[:-2]
+ shape = shape[-2:]
+ assert shape[0] == shape[1]
+ n = shape[0]
+ matrix = np.ones(shape).astype(np.float32) / (2.0 * n) + np.diag(
+ np.ones(n).astype(np.float32))
+ return np.tile(matrix, batch_shape + (1, 1))
+
+ def benchmarkLuOp(self):
+ for shape in self.shapes:
+ with ops.Graph().as_default(), \
+ session.Session(config=benchmark.benchmark_config()) as sess, \
+ ops.device("/cpu:0"):
+ matrix = variables.Variable(self._GenerateMatrix(shape))
+ lu, p = linalg_ops.lu(matrix)
+ variables.global_variables_initializer().run()
+ self.run_op_benchmark(
+ sess,
+ control_flow_ops.group(lu, p),
+ min_iters=25,
+ name="lu_cpu_{shape}".format(shape=shape))
+
+ if test.is_gpu_available(True):
+ with ops.Graph().as_default(), \
+ session.Session(config=benchmark.benchmark_config()) as sess, \
+ ops.device("/device:GPU:0"):
+ matrix = variables.Variable(self._GenerateMatrix(shape))
+ lu, p = linalg_ops.lu(matrix)
+ variables.global_variables_initializer().run()
+ self.run_op_benchmark(
+ sess,
+ control_flow_ops.group(lu, p),
+ min_iters=25,
+ name="lu_gpu_{shape}".format(shape=shape))
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/kernel_tests/matmul_op_test.py b/tensorflow/python/kernel_tests/matmul_op_test.py
index cd99f17..d31ecbc 100644
--- a/tensorflow/python/kernel_tests/matmul_op_test.py
+++ b/tensorflow/python/kernel_tests/matmul_op_test.py
@@ -21,6 +21,7 @@
import operator
import numpy as np
+from tensorflow.python import tf2
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
@@ -224,7 +225,8 @@
if __name__ == "__main__":
sizes = [1, 3, 5]
trans_options = [[False, False], [True, False], [False, True]]
- for use_static_shape in [False, True]:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_static_shape in set([True, tf2.enabled()]):
for dtype in (np.int32, np.int64, np.float16, np.float32, np.float64,
np.complex64, np.complex128):
if not use_static_shape and (dtype == np.int32 or dtype == np.int64):
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 129ea40..fdb7e4a 100644
--- a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py
+++ b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py
@@ -23,6 +23,7 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes as dtypes_lib
from tensorflow.python.framework import ops
+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 gradient_checker
@@ -44,6 +45,7 @@
def _GetMatrixBandPartTest(dtype_, batch_shape_, shape_):
+ @test_util.run_v1_only("b/120545219")
def Test(self):
mat = np.ones(shape_).astype(dtype_)
batch_mat = np.tile(mat, batch_shape_ + (1, 1))
@@ -73,6 +75,7 @@
def _GetMatrixBandPartGradTest(dtype_, batch_shape_, shape_):
+ @test_util.run_v1_only("b/120545219")
def Test(self):
shape = batch_shape_ + shape_
x = constant_op.constant(np.random.rand(*shape), dtype=dtype_)
diff --git a/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py b/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py
index b0bce6a..682ac12 100644
--- a/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py
+++ b/tensorflow/python/kernel_tests/matrix_logarithm_op_test.py
@@ -84,6 +84,7 @@
# Complex batch
self._verifyLogarithmComplex(self._makeBatch(matrix1, matrix2))
+ @test_util.run_v1_only("b/120545219")
def testNonSquareMatrix(self):
# When the logarithm of a non-square matrix is attempted we should return
# an error
@@ -91,6 +92,7 @@
gen_linalg_ops.matrix_logarithm(
np.array([[1., 2., 3.], [3., 4., 5.]], dtype=np.complex64))
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
# The input to the logarithm should be at least a 2-dimensional tensor.
tensor3 = constant_op.constant([1., 2.], dtype=dtypes.complex64)
@@ -121,6 +123,7 @@
size=np.prod(shape)).reshape(shape).astype(np.complex128)
self._verifyLogarithmComplex(matrix)
+ @test_util.run_v1_only("b/120545219")
def testConcurrentExecutesWithoutError(self):
with self.session(use_gpu=True) as sess:
matrix1 = math_ops.cast(
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 a6f5da9..463477a 100644
--- a/tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py
+++ b/tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py
@@ -20,6 +20,7 @@
import numpy as np
+from tensorflow.python import tf2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -134,7 +135,7 @@
self.assertEqual(np_ans.shape, tf_ans_val.shape)
self.assertAllClose(np_ans, tf_ans_val, atol=2 * tol, rtol=2 * tol)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
# The matrix and right-hand sides should have the same number of rows.
with self.session(use_gpu=True):
@@ -143,23 +144,26 @@
with self.assertRaises(ValueError):
linalg_ops.matrix_solve_ls(matrix, rhs)
- @test_util.run_deprecated_v1
def testEmpty(self):
full = np.array([[1., 2.], [3., 4.], [5., 6.]])
empty0 = np.empty([3, 0])
empty1 = np.empty([0, 2])
for fast in [True, False]:
with self.cached_session(use_gpu=True):
- tf_ans = linalg_ops.matrix_solve_ls(empty0, empty0, fast=fast).eval()
+ tf_ans = self.evaluate(
+ linalg_ops.matrix_solve_ls(empty0, empty0, fast=fast))
self.assertEqual(tf_ans.shape, (0, 0))
- tf_ans = linalg_ops.matrix_solve_ls(empty0, full, fast=fast).eval()
+ tf_ans = self.evaluate(
+ linalg_ops.matrix_solve_ls(empty0, full, fast=fast))
self.assertEqual(tf_ans.shape, (0, 2))
- tf_ans = linalg_ops.matrix_solve_ls(full, empty0, fast=fast).eval()
+ tf_ans = self.evaluate(
+ linalg_ops.matrix_solve_ls(full, empty0, fast=fast))
self.assertEqual(tf_ans.shape, (2, 0))
- tf_ans = linalg_ops.matrix_solve_ls(empty1, empty1, fast=fast).eval()
+ tf_ans = self.evaluate(
+ linalg_ops.matrix_solve_ls(empty1, empty1, fast=fast))
self.assertEqual(tf_ans.shape, (2, 2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBatchResultSize(self):
# 3x3x3 matrices, 3x3x1 right-hand sides.
matrix = np.array([1., 2., 3., 4., 5., 6., 7., 8., 9.] * 3).reshape(3, 3, 3)
@@ -350,7 +354,8 @@
if __name__ == "__main__":
for dtype_ in [np.float32, np.float64, np.complex64, np.complex128]:
- for use_placeholder_ in [True, False]:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_placeholder_ in set([False, not tf2.enabled()]):
for fast_ in [True, False]:
l2_regularizers = [0] if dtype_ == np.complex128 else [0, 0.1]
for l2_regularizer_ in l2_regularizers:
diff --git a/tensorflow/python/kernel_tests/matrix_square_root_op_test.py b/tensorflow/python/kernel_tests/matrix_square_root_op_test.py
index 1e2109b..3edb390 100644
--- a/tensorflow/python/kernel_tests/matrix_square_root_op_test.py
+++ b/tensorflow/python/kernel_tests/matrix_square_root_op_test.py
@@ -90,17 +90,20 @@
self._verifySquareRootReal(np.empty([0, 2, 2]))
self._verifySquareRootReal(np.empty([2, 0, 0]))
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
# The input to the square root should be at least a 2-dimensional tensor.
tensor = constant_op.constant([1., 2.])
with self.assertRaises(ValueError):
gen_linalg_ops.matrix_square_root(tensor)
+ @test_util.run_v1_only("b/120545219")
def testNotSquare(self):
with self.assertRaises(ValueError):
tensor = constant_op.constant([[1., 0., -1.], [-1., 1., 0.]])
self.evaluate(gen_linalg_ops.matrix_square_root(tensor))
+ @test_util.run_v1_only("b/120545219")
def testConcurrentExecutesWithoutError(self):
with test_util.use_gpu():
matrix1 = random_ops.random_normal([5, 5], seed=42)
diff --git a/tensorflow/python/kernel_tests/norm_op_test.py b/tensorflow/python/kernel_tests/norm_op_test.py
index 5ff0c58..20b9ad9 100644
--- a/tensorflow/python/kernel_tests/norm_op_test.py
+++ b/tensorflow/python/kernel_tests/norm_op_test.py
@@ -21,6 +21,7 @@
import numpy as np
from tensorflow.python.framework import constant_op
+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 as test_lib
@@ -35,6 +36,7 @@
class NormOpTest(test_lib.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testBadOrder(self):
matrix = [[0., 1.], [2., 3.]]
for ord_ in "fro", -7, -1.1, 0:
@@ -52,6 +54,7 @@
"'ord' must be a supported matrix norm"):
linalg_ops.norm(matrix, ord=ord_, axis=[-2, -1])
+ @test_util.run_v1_only("b/120545219")
def testInvalidAxis(self):
matrix = [[0., 1.], [2., 3.]]
for axis_ in [], [1, 2, 3], [[1]], [[1], [2]], [3.1415], [1, 1]:
@@ -78,6 +81,7 @@
tf_norm_val = sess.run(tf_norm, feed_dict={tf_matrix: matrix})
self.assertAllClose(np_norm, tf_norm_val, rtol=1e-5, atol=1e-5)
+ @test_util.run_v1_only("b/120545219")
def Test(self):
is_matrix_norm = (isinstance(axis_, tuple) or
isinstance(axis_, list)) and len(axis_) == 2
diff --git a/tensorflow/python/kernel_tests/numerics_test.py b/tensorflow/python/kernel_tests/numerics_test.py
index 5751f3f..f13f9d6 100644
--- a/tensorflow/python/kernel_tests/numerics_test.py
+++ b/tensorflow/python/kernel_tests/numerics_test.py
@@ -64,9 +64,9 @@
self.evaluate(t_verified)
+@test_util.run_v1_only("b/120545219")
class NumericsTest(test.TestCase):
- @test_util.run_deprecated_v1
def testInf(self):
with self.session(graph=ops.Graph()):
t1 = constant_op.constant(1.0)
@@ -77,7 +77,6 @@
with self.assertRaisesOpError("Inf"):
self.evaluate(a)
- @test_util.run_deprecated_v1
def testNaN(self):
with self.session(graph=ops.Graph()):
t1 = constant_op.constant(0.0)
@@ -88,7 +87,6 @@
with self.assertRaisesOpError("NaN"):
self.evaluate(a)
- @test_util.run_deprecated_v1
def testBoth(self):
with self.session(graph=ops.Graph()):
t1 = constant_op.constant([1.0, 0.0])
@@ -107,7 +105,6 @@
self.assertAllEqual(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), value)
self.assertEqual([2, 3], checked.get_shape())
- @test_util.run_deprecated_v1
def testControlFlowCond(self):
predicate = array_ops.placeholder(dtypes.bool, shape=[])
_ = control_flow_ops.cond(predicate,
@@ -120,7 +117,6 @@
r"or `tf.while_loop\(\)`\."):
numerics.add_check_numerics_ops()
- @test_util.run_deprecated_v1
def testControlFlowWhile(self):
predicate = array_ops.placeholder(dtypes.bool, shape=[])
_ = control_flow_ops.while_loop(lambda _: predicate,
diff --git a/tensorflow/python/kernel_tests/padding_fifo_queue_test.py b/tensorflow/python/kernel_tests/padding_fifo_queue_test.py
index b481836..e399969 100644
--- a/tensorflow/python/kernel_tests/padding_fifo_queue_test.py
+++ b/tensorflow/python/kernel_tests/padding_fifo_queue_test.py
@@ -29,11 +29,13 @@
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.platform import test
+@test_util.run_v1_only("b/120545219")
class PaddingFIFOQueueTest(test.TestCase):
def testConstructor(self):
@@ -1393,6 +1395,7 @@
with self.assertRaisesOpError("was cancelled"):
self.evaluate(enqueue_many_op)
+ @test_util.run_deprecated_v1
def testResetOfBlockingOperation(self):
with self.cached_session() as sess:
q_empty = data_flow_ops.PaddingFIFOQueue(5, dtypes_lib.float32, ((),))
diff --git a/tensorflow/python/kernel_tests/partitioned_variables_test.py b/tensorflow/python/kernel_tests/partitioned_variables_test.py
index 4865539..edcbc29 100644
--- a/tensorflow/python/kernel_tests/partitioned_variables_test.py
+++ b/tensorflow/python/kernel_tests/partitioned_variables_test.py
@@ -323,26 +323,24 @@
for i in xrange(len(expected_specs)):
self.assertEquals(expected_specs[i], slices[i]._save_slice_info.spec)
- @test_util.run_deprecated_v1
def testVecConstantInit(self):
with self.cached_session():
rnd_par = constant_op.constant([1, 2, 3, 4])
vs = partitioned_variables.create_partitioned_variables([4], [4], rnd_par)
- variables.global_variables_initializer().run()
- val = array_ops.concat(vs, 0).eval()
+ self.evaluate(variables.global_variables_initializer())
+ val = array_ops.concat(vs, 0)
rnd = self.evaluate(rnd_par)
self.assertAllClose(rnd, val)
self.assertEqual([dtypes.int32] * 4, [v.dtype.base_dtype for v in vs])
self._TestSaveSpec(vs, ["4 0,1", "4 1,1", "4 2,1", "4 3,1"])
- @test_util.run_deprecated_v1
def testConstantInit(self):
with self.cached_session():
rnd_par = constant_op.constant([[1, 2, 3, 4], [5, 6, 7, 8]])
vs = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
rnd_par)
- variables.global_variables_initializer().run()
- val = array_ops.concat(vs, 1).eval()
+ self.evaluate(variables.global_variables_initializer())
+ val = array_ops.concat(vs, 1)
rnd = self.evaluate(rnd_par)
self.assertAllClose(rnd, val)
self.assertEqual([dtypes.int32] * 2, [v.dtype.base_dtype for v in vs])
@@ -356,7 +354,7 @@
rnd_par)
vs2 = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
rnd_par)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
var1_name = vs1[0]._save_slice_info.full_name
var2_name = vs2[0]._save_slice_info.full_name
self.assertEqual("hi/PartitionedVariable", var1_name)
@@ -376,7 +374,7 @@
vs, reuse=True, use_resource=use_resource):
vs2 = partitioned_variables.create_partitioned_variables(
[2, 4], [1, 2], rnd_par, dtype=dtypes.int32)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
var1_name = vs1[0]._save_slice_info.full_name
var2_name = vs2[0]._save_slice_info.full_name
self.assertEqual("hola/PartitionedVariable", var1_name)
@@ -393,7 +391,7 @@
rnd_par)
vs2 = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
rnd_par)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
var1_name = vs1[0]._save_slice_info.full_name
var2_name = vs2[0]._save_slice_info.full_name
# Currently, the name scope 'ola' has no effect.
@@ -408,18 +406,16 @@
def testName(self):
self._testNameHelper(use_resource=False)
- @test_util.run_deprecated_v1
def testResourceName(self):
self._testNameHelper(use_resource=True)
- @test_util.run_deprecated_v1
def testRandomInitValue(self):
with self.cached_session():
rnd = variables.Variable(random_ops.random_uniform([200, 40]))
vs = partitioned_variables.create_partitioned_variables(
rnd.get_shape(), [1, 10], rnd.initialized_value())
- variables.global_variables_initializer().run()
- val = array_ops.concat(vs, 1).eval()
+ self.evaluate(variables.global_variables_initializer())
+ val = array_ops.concat(vs, 1)
rnd = self.evaluate(rnd)
self.assertAllClose(rnd, val)
self.assertEqual([dtypes.float32] * 10, [v.dtype.base_dtype for v in vs])
@@ -430,7 +426,6 @@
"200 40 0,200:36,4"
])
- @test_util.run_deprecated_v1
def testRandomInitUnevenPartitions(self):
with self.cached_session():
rnd = variables.Variable(
@@ -440,7 +435,7 @@
rnd.get_shape(), [1, i], rnd.initialized_value())
for i in xrange(1, 10)
]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
rnd_val = self.evaluate(rnd)
# Only check the slice save specs for the first 5 tf.
save_specs = [
@@ -462,33 +457,31 @@
]
]
for i, vs in enumerate(var_lists):
- var_val = array_ops.concat(vs, 1).eval()
+ var_val = array_ops.concat(vs, 1)
self.assertAllClose(rnd_val, var_val)
self.assertEqual([dtypes.float64] * len(vs),
[v.dtype.base_dtype for v in vs])
if i < len(save_specs):
self._TestSaveSpec(vs, save_specs[i])
- @test_util.run_deprecated_v1
def testDegenerate(self):
with self.cached_session():
rnd = variables.Variable(random_ops.random_uniform([10, 43]))
vs = partitioned_variables.create_partitioned_variables(
rnd.get_shape(), [1, 1], rnd.initialized_value())
- variables.global_variables_initializer().run()
- val = array_ops.concat(vs, 0).eval()
+ self.evaluate(variables.global_variables_initializer())
+ val = array_ops.concat(vs, 0)
rnd = self.evaluate(rnd)
self.assertAllClose(rnd, val)
self._TestSaveSpec(vs, ["10 43 0,10:0,43"])
- @test_util.run_deprecated_v1
def testSliceSizeOne(self):
with self.cached_session():
rnd = variables.Variable(random_ops.random_uniform([10, 43]))
vs = partitioned_variables.create_partitioned_variables(
rnd.get_shape(), [10, 1], rnd.initialized_value())
- variables.global_variables_initializer().run()
- val = array_ops.concat(vs, 0).eval()
+ self.evaluate(variables.global_variables_initializer())
+ val = array_ops.concat(vs, 0)
rnd = self.evaluate(rnd)
self.assertAllClose(rnd, val)
self._TestSaveSpec(vs, [
@@ -497,7 +490,6 @@
"10 43 6,1:0,43", "10 43 7,1:0,43", "10 43 8,1:0,43", "10 43 9,1:0,43"
])
- @test_util.run_deprecated_v1
def testIotaInitializer(self):
self.assertAllClose([0., 1., 2., 3.], _IotaInitializer([4]))
self.assertAllClose([[0., 1.], [0., 10.], [0., 100.], [0., 1000.]],
@@ -505,11 +497,11 @@
with self.cached_session():
vs = partitioned_variables.create_partitioned_variables([13, 5], [3, 1],
_IotaInitializer)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
slice0 = _IotaInitializer([5, 5])
slice1 = _IotaInitializer([4, 5])
slice2 = _IotaInitializer([4, 5])
- val = array_ops.concat(vs, 0).eval()
+ val = array_ops.concat(vs, 0)
self.assertAllClose(slice0 + slice1 + slice2, val)
self._TestSaveSpec(vs, ["13 5 0,5:0,5", "13 5 5,4:0,5", "13 5 9,4:0,5"])
@@ -520,7 +512,7 @@
with self.cached_session():
var0, var1 = partitioned_variables.create_partitioned_variables(
[20, 12], [1, 2], init_ops.random_uniform_initializer())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
val0, val1 = self.evaluate(var0).flatten(), self.evaluate(var1).flatten()
self.assertTrue(np.linalg.norm(val0 - val1) > 1e-6)
# Negative test that proves that slices have the same values if
@@ -528,7 +520,7 @@
with self.cached_session():
var0, var1 = partitioned_variables.create_partitioned_variables(
[20, 12], [1, 2], init_ops.random_uniform_initializer(seed=201))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
val0, val1 = self.evaluate(var0).flatten(), self.evaluate(var1).flatten()
self.assertAllClose(val0, val1)
@@ -607,8 +599,8 @@
self.assertTrue(
c.op in concat_control_inputs,
"var_x._concat() should get control dependencies from its scope.")
- variables.global_variables_initializer().run()
- self.assertAllClose(value.eval(), var_x.as_tensor().eval())
+ self.evaluate(variables.global_variables_initializer())
+ self.assertAllClose(value, var_x.as_tensor())
def testMetaGraphSaveLoad(self):
save_prefix = os.path.join(self.get_temp_dir(), "ckpt")
@@ -623,7 +615,7 @@
v0_part = v0._get_partitions()
self.assertEqual(len(v0_list), 5)
self.assertAllEqual(v0_part, (5, 1))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
save_graph.get_collection_ref("partvar").append(v0)
saver = saver_lib.Saver()
diff --git a/tensorflow/python/kernel_tests/priority_queue_test.py b/tensorflow/python/kernel_tests/priority_queue_test.py
index 9be682e..49ec7ee 100644
--- a/tensorflow/python/kernel_tests/priority_queue_test.py
+++ b/tensorflow/python/kernel_tests/priority_queue_test.py
@@ -27,6 +27,7 @@
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 test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
import tensorflow.python.ops.nn_grad # pylint: disable=unused-import
@@ -35,6 +36,7 @@
class PriorityQueueTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testRoundTripInsertReadOnceSorts(self):
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.string, dtypes.string), (
@@ -112,6 +114,7 @@
missed.remove((dv0, dv1))
self.assertEqual(missed, set())
+ @test_util.run_v1_only("b/120545219")
def testRoundTripFillsCapacityMultiThreadedEnqueueAndDequeue(self):
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(10, (dtypes.int64), (()))
@@ -267,6 +270,7 @@
missed.remove((dv0, dv1))
self.assertEqual(missed, set())
+ @test_util.run_v1_only("b/120545219")
def testRoundTripInsertOnceReadOnceSorts(self):
with self.cached_session() as sess:
q = data_flow_ops.PriorityQueue(2000, (dtypes.string, dtypes.string), (
@@ -288,6 +292,7 @@
for e, dv0, dv1 in zip(deq_elem, deq_value_0, deq_value_1):
self.assertTrue((dv0, dv1) in allowed[e])
+ @test_util.run_v1_only("b/120545219")
def testRoundTripInsertOnceReadManySorts(self):
with self.cached_session():
q = data_flow_ops.PriorityQueue(2000, (dtypes.int64), (()))
@@ -296,6 +301,7 @@
deq_values = np.hstack((q.dequeue_many(100)[0].eval() for _ in range(10)))
self.assertAllEqual(deq_values, sorted(elem))
+ @test_util.run_v1_only("b/120545219")
def testRoundTripInsertOnceReadOnceLotsSorts(self):
with self.cached_session():
q = data_flow_ops.PriorityQueue(2000, (dtypes.int64), (()))
@@ -311,6 +317,7 @@
with self.assertRaises(TypeError):
q.enqueue_many((["a", "b", "c"], ["a", "b", "c"])).run()
+ @test_util.run_v1_only("b/120545219")
def testInsertingNonScalarFails(self):
with self.cached_session() as sess:
input_priority = array_ops.placeholder(dtypes.int64)
diff --git a/tensorflow/python/kernel_tests/py_func_test.py b/tensorflow/python/kernel_tests/py_func_test.py
index 1f3f02a..482633d 100644
--- a/tensorflow/python/kernel_tests/py_func_test.py
+++ b/tensorflow/python/kernel_tests/py_func_test.py
@@ -102,6 +102,7 @@
script_ops.eager_py_func(np_func, [x, y], [dtypes.float32]))
self.assertEqual(z[0], np_func(1.0, 2.0).astype(np.float32))
+ @test_util.run_v1_only("b/120545219")
def testArray(self):
with self.cached_session():
x = constant_op.constant([1.0, 2.0], dtypes.float64)
@@ -168,6 +169,7 @@
(dtypes.float64, dtypes.float64)))
self.assertAllClose(y, [0.0, 1.0])
+ @test_util.run_v1_only("b/120545219")
def testStrings(self):
def read_fixed_length_numpy_strings():
@@ -185,6 +187,7 @@
script_ops.py_func(read_and_return_strings, [x, y], dtypes.string))
self.assertAllEqual(z, [b"hello there", b"hi there"])
+ @test_util.run_v1_only("b/120545219")
def testStringsAreConvertedToBytes(self):
def read_fixed_length_numpy_strings():
@@ -202,6 +205,7 @@
script_ops.py_func(read_and_return_strings, [x, y], dtypes.string))
self.assertAllEqual(z, [b"hello there", b"hi there"])
+ @test_util.run_v1_only("b/120545219")
def testObjectArraysAreConvertedToBytes(self):
def read_object_array():
@@ -217,12 +221,14 @@
z, = script_ops.py_func(read_and_return_strings, [x, y], [dtypes.string])
self.assertListEqual(list(z.eval()), [b"hello there", b"hi ya"])
+ @test_util.run_v1_only("b/120545219")
def testStringPadding(self):
correct = [b"this", b"is", b"a", b"test"]
with self.cached_session():
s, = script_ops.py_func(lambda: [correct], [], [dtypes.string])
self.assertAllEqual(s.eval(), correct)
+ @test_util.run_v1_only("b/120545219")
def testStringPaddingAreConvertedToBytes(self):
inp = ["this", "is", "a", "test"]
correct = [b"this", b"is", b"a", b"test"]
@@ -230,6 +236,7 @@
s, = script_ops.py_func(lambda: [inp], [], [dtypes.string])
self.assertAllEqual(s.eval(), correct)
+ @test_util.run_v1_only("b/120545219")
def testLarge(self):
with self.cached_session() as sess:
x = array_ops.zeros([1000000], dtype=np.float32)
@@ -243,6 +250,7 @@
x = self.evaluate(script_ops.py_func(lambda: 42.0, [], dtypes.float64))
self.assertAllClose(x, 42.0)
+ @test_util.run_v1_only("b/120545219")
def testAlias(self):
with self.cached_session():
np_array = np.array([1.0, 2.0], dtype=np.float32)
@@ -251,6 +259,7 @@
value.op.run()
self.assertAllEqual(np_array, [1.0, 2.0])
+ @test_util.run_v1_only("b/120545219")
def testReturnUnicodeString(self):
with self.cached_session():
correct = u"你好 世界"
@@ -261,6 +270,7 @@
z, = script_ops.py_func(unicode_string, [], [dtypes.string])
self.assertEqual(z.eval(), correct.encode("utf8"))
+ @test_util.run_v1_only("b/120545219")
def testBadNumpyReturnType(self):
with self.cached_session():
@@ -274,6 +284,7 @@
"Unsupported numpy type"):
self.evaluate(y)
+ @test_util.run_v1_only("b/120545219")
def testBadReturnType(self):
with self.cached_session():
@@ -287,6 +298,7 @@
"Unsupported object type"):
self.evaluate(z)
+ @test_util.run_v1_only("b/120545219")
def testReturnInput(self):
with self.cached_session():
@@ -321,6 +333,7 @@
self.assertEqual(self.evaluate(x), 0)
self.assertEqual(self.evaluate(x), 0)
+ @test_util.run_v1_only("b/120545219")
def testGradientFunction(self):
# Input to tf.py_func is necessary, otherwise get_gradient_function()
# returns None per default.
@@ -330,6 +343,7 @@
self.assertEqual(None, ops.get_gradient_function(x.op))
self.assertEqual(None, ops.get_gradient_function(y.op))
+ @test_util.run_v1_only("b/120545219")
def testCOrder(self):
with self.cached_session():
val = [[1, 2], [3, 4]]
@@ -337,6 +351,7 @@
[dtypes.int64])
self.assertAllEqual(val, self.evaluate(x))
+ @test_util.run_v1_only("b/120545219")
def testParallel(self):
# Tests that tf.py_func's can run in parallel if they release the GIL.
with self.cached_session() as session:
@@ -382,6 +397,7 @@
self.assertIsNone(ret)
self.assertAllEqual([3], s.value)
+ @test_util.run_v1_only("b/120545219")
def testNoReturnValueStateless(self):
def do_nothing(unused_x):
@@ -420,6 +436,7 @@
with self.assertRaisesWithPredicateMatch(tf_exp, expected_error_check):
self.evaluate(f)
+ @test_util.run_v1_only("b/120545219")
def testExceptionHandling(self):
with self.cached_session():
self._testExceptionHandling(ValueError, errors.InvalidArgumentError)
@@ -514,7 +531,7 @@
self.assertAllEqual(ret, [[3.0], [3.0], [3.0]])
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testEagerExceptionHandling(self):
with test_util.device(use_gpu=True):
self._testExceptionHandling(
@@ -534,7 +551,7 @@
self._testExceptionHandling(WeirdError, errors.UnknownError, eager=True)
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testEagerReturningVariableRaisesError(self):
def return_variable():
return resource_variable_ops.ResourceVariable(0.0)
@@ -558,6 +575,7 @@
dy_dx = tape.gradient(y, x)
self.assertEqual(self.evaluate(dy_dx), 6.0)
+ @test_util.run_v1_only("b/120545219")
def testEagerGradientGraph(self):
def f(x):
@@ -568,6 +586,7 @@
dy_dx = gradients_impl.gradients(y, x)[0]
self.assertEqual(self.evaluate(dy_dx), 6.0)
+ @test_util.run_v1_only("b/120545219")
def testEagerGradientGraphTwoOutputs(self):
def f(x, y):
@@ -597,6 +616,7 @@
self.assertEqual(self.evaluate(dz_dx), 6.0)
self.assertEqual(self.evaluate(dz_dy), 8.0)
+ @test_util.run_v1_only("b/120545219")
def testEagerGradientGraphMultipleArgs(self):
def f(x, y):
@@ -610,6 +630,7 @@
self.assertEqual(self.evaluate(dz_dx), 6.0)
self.assertEqual(self.evaluate(dz_dy), 8.0)
+ @test_util.run_v1_only("b/120545219")
def testEagerGradientGraphLogHuber(self):
def log_huber(x, m):
@@ -631,6 +652,7 @@
self.assertEqual(y, 1.0)
self.assertEqual(dy_dx, 2.0)
+ @test_util.run_v1_only("b/120545219")
def testEagerRespectsDevicePlacmentOfOp(self):
def f(x):
diff --git a/tensorflow/python/kernel_tests/qr_op_test.py b/tensorflow/python/kernel_tests/qr_op_test.py
index 0f2537b..5adb95c 100644
--- a/tensorflow/python/kernel_tests/qr_op_test.py
+++ b/tensorflow/python/kernel_tests/qr_op_test.py
@@ -20,6 +20,7 @@
import numpy as np
+from tensorflow.python import tf2
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
@@ -39,6 +40,7 @@
class QrOpTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
# The input to qr should be a tensor of at least rank 2.
scalar = constant_op.constant(1.)
@@ -102,7 +104,7 @@
tol = 1e-14
# Tests that a ~= q*r.
a_recon = math_ops.matmul(q, r)
- self.assertAllClose(a_recon.eval(), a, rtol=tol, atol=tol)
+ self.assertAllClose(a_recon, a, rtol=tol, atol=tol)
def CheckUnitary(self, x):
# Tests that x[...,:,:]^H * x[...,:,:] is close to the identity.
@@ -112,8 +114,9 @@
tol = 1e-5
else:
tol = 1e-14
- self.assertAllClose(identity.eval(), self.evaluate(xx), atol=tol)
+ self.assertAllClose(identity, xx, atol=tol)
+ @test_util.run_v1_only("b/120545219")
def Test(self):
np.random.seed(1)
x_np = np.random.uniform(
@@ -162,6 +165,7 @@
def _GetQrGradOpTest(dtype_, shape_, full_matrices_):
+ @test_util.run_v1_only("b/120545219")
def Test(self):
np.random.seed(42)
a = np.random.uniform(low=-1.0, high=1.0, size=shape_).astype(dtype_)
@@ -202,7 +206,8 @@
for cols in 1, 2, 5, 10, 32, 100:
for full_matrices in False, True:
for batch_dims in [(), (3,)] + [(3, 2)] * (max(rows, cols) < 10):
- for use_static_shape in True, False:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_static_shape in set([True, tf2.enabled()]):
shape = batch_dims + (rows, cols)
name = "%s_%s_full_%s_static_%s" % (dtype.__name__,
"_".join(map(str, shape)),
diff --git a/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py b/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py
index ed4f543..dd814a2 100644
--- a/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py
+++ b/tensorflow/python/kernel_tests/random/random_shuffle_queue_test.py
@@ -29,11 +29,13 @@
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging
+@test_util.run_v1_only("b/120545219")
class RandomShuffleQueueTest(test.TestCase):
def setUp(self):
@@ -1415,6 +1417,7 @@
self.assertItemsEqual(elem, results)
+ @test_util.run_v1_only("b/120545219")
def testBigDequeueMany(self):
with self.cached_session() as sess:
q = data_flow_ops.RandomShuffleQueue(2, 0, dtypes_lib.int32, ((),))
diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
index b57d9d4..df7b686 100644
--- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py
+++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
@@ -33,7 +33,10 @@
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 custom_gradient
+from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import list_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import state_ops
@@ -585,6 +588,33 @@
v.load(2.0)
self.assertEqual(2.0, self.evaluate(v.value()))
+ def testShapePassedToGradient(self):
+ with ops.Graph().as_default():
+ @custom_gradient.custom_gradient
+ def differentiable_scatter_update(handle, indices, values):
+ with ops.control_dependencies([
+ resource_variable_ops.resource_scatter_update(
+ handle, indices, values)]):
+ new_handle = array_ops.identity(handle)
+
+ def grad(dresult):
+ self.assertIsNotNone(
+ tensor_util.constant_value(dresult.dense_shape))
+ return [dresult, None, None]
+
+ return new_handle, grad
+
+ var = variable_scope.get_variable(
+ "foo", shape=[20], initializer=init_ops.zeros_initializer,
+ dtype=dtypes.float64, use_resource=True)
+
+ indices = math_ops.range(10)
+ updates = math_ops.range(9, -1, -1, dtype=dtypes.float64)
+ new_handle = differentiable_scatter_update(var.handle, indices, updates)
+ gathered = resource_variable_ops.resource_gather(
+ new_handle, indices, dtype=var.dtype)
+ gradients_impl.gradients([gathered], [updates])
+
def testToFromProtoCachedValue(self):
with ops.Graph().as_default():
v_def = resource_variable_ops.ResourceVariable(
@@ -599,7 +629,7 @@
variable_def=other_v_def)
self.assertTrue(other_v_prime._cached_value is not None)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testVariableDefInitializedInstances(self):
with ops.Graph().as_default(), self.cached_session() as sess:
v_def = resource_variable_ops.ResourceVariable(
@@ -659,7 +689,7 @@
def testToFromProto(self):
with self.cached_session():
v = resource_variable_ops.ResourceVariable(1.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
w = resource_variable_ops.ResourceVariable.from_proto(v.to_proto())
self.assertEquals(2, math_ops.add(w, 1).eval())
@@ -704,7 +734,7 @@
self.assertEqual(0.0, self.evaluate(v.value()))
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testDestroyResource(self):
v = resource_variable_ops.ResourceVariable(3.0, name="var0")
self.evaluate(variables.global_variables_initializer())
@@ -767,7 +797,7 @@
def testSharedName(self):
with self.cached_session():
v = resource_variable_ops.ResourceVariable(300.0, name="var4")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
w = resource_variable_ops.var_handle_op(
dtype=v.dtype.base_dtype, shape=v.get_shape(), shared_name="var4",
@@ -820,7 +850,7 @@
v.initializer.run(feed_dict={v.initial_value: 3.0})
self.assertEqual(3.0, v.value().eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testControlFlowInitialization(self):
"""Expects an error if an initializer is in a control-flow scope."""
@@ -924,6 +954,19 @@
state_ops.scatter_sub(v, [1], [3])
self.assertAllEqual([1.0, -1.0], v.numpy())
+ def testScatterUpdateVariant(self):
+ with context.eager_mode():
+ v = resource_variable_ops.ResourceVariable([
+ list_ops.empty_tensor_list(
+ element_dtype=dtypes.float32, element_shape=[])
+ ])
+ v.scatter_update(
+ ops.IndexedSlices(
+ list_ops.tensor_list_from_tensor([1., 2.], element_shape=[]), 0))
+ self.assertAllEqual(
+ list_ops.tensor_list_get_item(v[0], 0, element_dtype=dtypes.float32),
+ 1.)
+
def testScatterNdAddStateOps(self):
with context.eager_mode():
v = resource_variable_ops.ResourceVariable(
@@ -957,7 +1000,7 @@
self.assertAllEqual(self.evaluate(v.assign_add(1)), [1, 2, 3, 4])
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCopyToGraphUninitialized(self):
v = resource_variable_ops.ResourceVariable([0, 1, 2, 3])
copy_to_graph = ops.Graph()
diff --git a/tensorflow/python/kernel_tests/rnn_test.py b/tensorflow/python/kernel_tests/rnn_test.py
index 3bc457f..a49496e 100644
--- a/tensorflow/python/kernel_tests/rnn_test.py
+++ b/tensorflow/python/kernel_tests/rnn_test.py
@@ -262,7 +262,7 @@
rnn.dynamic_rnn(cell, inputs, dtype=dtypes.float32, sequence_length=[4])
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testTensorArrayStateIsAccepted(self):
cell = TensorArrayStateRNNCell()
in_eager_mode = context.executing_eagerly()
diff --git a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
index c1241ba..8510a08 100644
--- a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
+++ b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
@@ -217,7 +217,7 @@
def testVariableRankAdd(self):
self._VariableRankTests(_NumpyAdd, state_ops.scatter_nd_add)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testVariableRankSub(self):
self._VariableRankTests(_NumpySub, state_ops.scatter_nd_sub)
@@ -235,7 +235,7 @@
self._VariableRankTest(
np_scatter, tf_scatter, vtype, itype, repeat_indices=True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testScatterRepeatIndices(self):
"""This tests scatter_add using indices that repeat."""
self._ScatterRepeatIndicesTest(_NumpyAdd, state_ops.scatter_nd_add)
@@ -257,7 +257,7 @@
# session.run([update0, update1])
# self.assertAllEqual([False, True], self.evaluate(var))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testScatterOutOfRangeCpu(self):
# TODO(simister): Re-enable once binary size increase due to
# scatter_nd ops is under control.
@@ -294,7 +294,7 @@
state_ops.scatter_nd_update(ref, indices,
updates).get_shape().as_list(), shape)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testResVarInvalidOutputShape(self):
res = variables.Variable(
initial_value=lambda: array_ops.zeros(shape=[], dtype=dtypes.float32),
@@ -509,7 +509,7 @@
ValueError, "Indices and updates specified for empty output shape"):
self.scatter_nd(indices, updates, shape)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testEmptyOutputShape2(self):
indices = array_ops.placeholder(dtypes.int32, shape=None)
updates = array_ops.placeholder(dtypes.int32, shape=None)
@@ -717,6 +717,7 @@
self.assertAllEqual(subbed,
constant_op.constant([1, -10, 1, -9, -8, 1, 1, -11]))
+ @test_util.run_v1_only("b/120545219")
def testUpdateAddSubGradients(self):
with self.cached_session():
diff --git a/tensorflow/python/kernel_tests/scatter_ops_test.py b/tensorflow/python/kernel_tests/scatter_ops_test.py
index 623c17d..ce7e0c0 100644
--- a/tensorflow/python/kernel_tests/scatter_ops_test.py
+++ b/tensorflow/python/kernel_tests/scatter_ops_test.py
@@ -192,6 +192,10 @@
if tf_scatter != state_ops.scatter_div:
vtypes.append(np.int32)
+ if (tf_scatter == state_ops.scatter_min or
+ tf_scatter == state_ops.scatter_max):
+ vtypes.append(np.float16)
+
for vtype in vtypes:
for itype in (np.int32, np.int64):
self._VariableRankTest(tf_scatter, vtype, itype, repeat_indices,
diff --git a/tensorflow/python/kernel_tests/self_adjoint_eig_op_test.py b/tensorflow/python/kernel_tests/self_adjoint_eig_op_test.py
index 42577f7..47b22ec 100644
--- a/tensorflow/python/kernel_tests/self_adjoint_eig_op_test.py
+++ b/tensorflow/python/kernel_tests/self_adjoint_eig_op_test.py
@@ -22,8 +22,9 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes as dtypes_lib
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
@@ -39,6 +40,7 @@
class SelfAdjointEigTest(test.TestCase):
+ @test_util.run_deprecated_v1
def testWrongDimensions(self):
# The input to self_adjoint_eig should be a tensor of
# at least rank 2.
@@ -49,6 +51,7 @@
with self.assertRaises(ValueError):
linalg_ops.self_adjoint_eig(vector)
+ @test_util.run_deprecated_v1
def testConcurrentExecutesWithoutError(self):
all_ops = []
with self.session(use_gpu=True) as sess:
@@ -161,7 +164,7 @@
math_ops.matmul(tf_v, array_ops.matrix_diag(tf_e)),
tf_v,
adjoint_b=True)
- self.assertAllClose(a_ev.eval(), a, atol=atol)
+ self.assertAllClose(self.evaluate(a_ev), a, atol=atol)
# Compare to numpy.linalg.eigh.
CompareEigenDecompositions(self, np_e, np_v, self.evaluate(tf_e),
@@ -169,7 +172,7 @@
else:
tf_e = linalg_ops.self_adjoint_eigvals(constant_op.constant(a))
self.assertAllClose(
- np.sort(np_e, -1), np.sort(tf_e.eval(), -1), atol=atol)
+ np.sort(np_e, -1), np.sort(self.evaluate(tf_e), -1), atol=atol)
return Test
@@ -185,53 +188,51 @@
n = shape_[-1]
batch_shape = shape_[:-2]
np_dtype = dtype_.as_numpy_dtype
- a = np.random.uniform(
- low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
- if dtype_.is_complex:
- a += 1j * np.random.uniform(
+
+ def RandomInput():
+ a = np.random.uniform(
low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
- a += np.conj(a.T)
- a = np.tile(a, batch_shape + (1, 1))
+ if dtype_.is_complex:
+ a += 1j * np.random.uniform(
+ low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
+ a += np.conj(a.T)
+ a = np.tile(a, batch_shape + (1, 1))
+ return a
+
# Optimal stepsize for central difference is O(epsilon^{1/3}).
epsilon = np.finfo(np_dtype).eps
delta = 0.1 * epsilon**(1.0 / 3.0)
# tolerance obtained by looking at actual differences using
# np.linalg.norm(theoretical-numerical, np.inf) on -mavx build
+ # after discarding one random input sample
+ _ = RandomInput()
if dtype_ in (dtypes_lib.float32, dtypes_lib.complex64):
tol = 1e-2
else:
tol = 1e-7
with self.session(use_gpu=True):
- tf_a = constant_op.constant(a)
- if compute_v_:
- tf_e, tf_v = linalg_ops.self_adjoint_eig(tf_a)
+ def Compute(x):
+ e, v = linalg_ops.self_adjoint_eig(x)
# (complex) Eigenvectors are only unique up to an arbitrary phase
# We normalize the vectors such that the first component has phase 0.
- top_rows = tf_v[..., 0:1, :]
- if tf_a.dtype.is_complex:
+ top_rows = v[..., 0:1, :]
+ if dtype_.is_complex:
angle = -math_ops.angle(top_rows)
phase = math_ops.complex(math_ops.cos(angle), math_ops.sin(angle))
else:
phase = math_ops.sign(top_rows)
- tf_v *= phase
- outputs = [tf_e, tf_v]
+ v *= phase
+ return e, v
+
+ if compute_v_:
+ funcs = [lambda x: Compute(x)[0], lambda x: Compute(x)[1]]
else:
- tf_e = linalg_ops.self_adjoint_eigvals(tf_a)
- outputs = [tf_e]
- for b in outputs:
- x_init = np.random.uniform(
- low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
- if dtype_.is_complex:
- x_init += 1j * np.random.uniform(
- low=-1.0, high=1.0, size=n * n).reshape([n, n]).astype(np_dtype)
- x_init += np.conj(x_init.T)
- x_init = np.tile(x_init, batch_shape + (1, 1))
- theoretical, numerical = gradient_checker.compute_gradient(
- tf_a,
- tf_a.get_shape().as_list(),
- b,
- b.get_shape().as_list(),
- x_init_value=x_init,
+ funcs = [linalg_ops.self_adjoint_eigvals]
+
+ for f in funcs:
+ theoretical, numerical = gradient_checker_v2.compute_gradient(
+ f,
+ [RandomInput()],
delta=delta)
self.assertAllClose(theoretical, numerical, atol=tol, rtol=tol)
@@ -245,7 +246,7 @@
for size in 1, 2, 5, 10:
for batch_dims in [(), (3,)] + [(3, 2)] * (max(size, size) < 10):
shape = batch_dims + (size, size)
- name = "%s_%s_%s" % (dtype, "_".join(map(str, shape)), compute_v)
+ name = "%s_%s_%s" % (dtype.name, "_".join(map(str, shape)), compute_v)
_AddTest(SelfAdjointEigTest, "SelfAdjointEig", name,
_GetSelfAdjointEigTest(dtype, shape, compute_v))
_AddTest(SelfAdjointEigGradTest, "SelfAdjointEigGrad", name,
diff --git a/tensorflow/python/kernel_tests/session_ops_test.py b/tensorflow/python/kernel_tests/session_ops_test.py
index dc663cb..bc5d8e8 100644
--- a/tensorflow/python/kernel_tests/session_ops_test.py
+++ b/tensorflow/python/kernel_tests/session_ops_test.py
@@ -20,6 +20,7 @@
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.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import session_ops
@@ -28,6 +29,7 @@
from tensorflow.python.platform import test
+@test_util.run_v1_only("b/120545219")
class SessionOpsTest(test.TestCase):
def testHandleBasic(self):
@@ -232,6 +234,7 @@
b_p: b_handle.handle})
self.assertEqual(3.0, c_handle.eval())
+ @test_util.run_v1_only("b/120545219")
def testFeedOneHandleDirectly(self):
with self.cached_session() as sess:
a = constant_op.constant(10.0)
@@ -243,6 +246,7 @@
self.assertAllClose(2500.0, sess.run(d, feed_dict={c: h_c}))
+ @test_util.run_v1_only("b/120545219")
def testDirectHandleFeedOverlappingWithFetches(self):
with self.cached_session() as sess:
a = constant_op.constant(10.0)
@@ -283,6 +287,7 @@
self.assertAllClose(48.0, sess.run(e, feed_dict={c: h_c, d: h_d}))
self.assertAllClose(-48.0, sess.run(e, feed_dict={c: h_d, d: h_c}))
+ @test_util.run_v1_only("b/120545219")
def testFeedHandleToVariableDirectly(self):
with self.cached_session() as sess:
a = variables.Variable(12.0)
diff --git a/tensorflow/python/kernel_tests/signal/reconstruction_ops_test.py b/tensorflow/python/kernel_tests/signal/reconstruction_ops_test.py
index 4cb6ced..e0ce064 100644
--- a/tensorflow/python/kernel_tests/signal/reconstruction_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/reconstruction_ops_test.py
@@ -66,6 +66,7 @@
self.assertAllClose(output, expected_output)
+ @test_util.run_deprecated_v1
def test_unknown_shapes(self):
# This test uses placeholders and does not work in eager mode.
if context.executing_eagerly():
@@ -85,6 +86,7 @@
self.assertAllClose(output, expected_output)
+ @test_util.run_deprecated_v1
def test_unknown_rank(self):
# This test uses placeholders and does not work in eager mode.
if context.executing_eagerly():
@@ -104,6 +106,7 @@
self.assertAllClose(output, expected_output)
+ @test_util.run_deprecated_v1
def test_fast_path(self):
# This test uses tensor names and does not work in eager mode.
if context.executing_eagerly():
diff --git a/tensorflow/python/kernel_tests/sparse_conditional_accumulator_test.py b/tensorflow/python/kernel_tests/sparse_conditional_accumulator_test.py
index 275c86e..4a967b6 100644
--- a/tensorflow/python/kernel_tests/sparse_conditional_accumulator_test.py
+++ b/tensorflow/python/kernel_tests/sparse_conditional_accumulator_test.py
@@ -267,7 +267,7 @@
self.assertAllEqual(val.values, [[5, 5], [0, 20], [30, 0]])
self.assertAllEqual(val.dense_shape, [-1, 2])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testParallelApplyGradMean(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -299,7 +299,7 @@
np.array([[expected_val, 0], [0, expected_val]]).astype(np.float32),
val, sess)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testParallelApplyGradSum(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -334,7 +334,7 @@
np.array([[expected_val, 0], [0, expected_val]]).astype(np.float32),
val, sess)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testParallelTakeGrad(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -374,7 +374,7 @@
self._assertEqual_nparray(
np.array([[0, 0], [elems[i], 0]]), results[i], sess)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAccumulatorApplyAndBlockingTake(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -410,7 +410,7 @@
with self.assertRaisesOpError("was cancelled"):
self.evaluate(takeg_op)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAccumulatorCancel(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -430,7 +430,7 @@
takeg_thread.join()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNonVectorIndices(self):
with self.cached_session():
q = data_flow_ops.SparseConditionalAccumulator(
@@ -443,7 +443,7 @@
grad_indices=[[0, 1], [1, 0]],
grad_values=np.array([1, 2]).astype(np.float32)).run()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testZeroDimensionValues(self):
with self.cached_session():
q = data_flow_ops.SparseConditionalAccumulator(
@@ -454,7 +454,7 @@
q.apply_grad(
grad_indices=[0], grad_values=np.array(1).astype(np.float32)).run()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWrongNonEmptyInputValues(self):
with self.cached_session():
q = data_flow_ops.SparseConditionalAccumulator(
@@ -466,7 +466,7 @@
grad_indices=[0, 1],
grad_values=np.array([[0, 1, 1]]).astype(np.float32)).run()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testDynamicNonVectorIndices(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -486,7 +486,7 @@
x_values: np.array([1, 2]).astype(np.float32)
})
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testDynamicWrongNonEmptyInputValues(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
@@ -505,7 +505,7 @@
x_values: np.array([[0, 1, 1]]).astype(np.float32)
})
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testEmptyShapeApply(self):
with self.cached_session():
q = data_flow_ops.SparseConditionalAccumulator(
@@ -531,7 +531,7 @@
q.apply_grad(grad_indices=[0], grad_values=[1.0], grad_shape=[]).run()
q.apply_grad(grad_indices=[0], grad_values=[1.0]).run()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testValidateShape(self):
with self.cached_session() as sess:
q = data_flow_ops.SparseConditionalAccumulator(
diff --git a/tensorflow/python/kernel_tests/sparse_ops_test.py b/tensorflow/python/kernel_tests/sparse_ops_test.py
index 75f65e6..7598991 100644
--- a/tensorflow/python/kernel_tests/sparse_ops_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops_test.py
@@ -22,6 +22,7 @@
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 sparse_tensor
from tensorflow.python.framework import test_util
@@ -799,6 +800,19 @@
self.assertAllEqual(result_np, res_densified)
@test_util.run_deprecated_v1
+ def testCwiseShapeValidation(self):
+ # Test case for GitHub 24072.
+ with self.session(use_gpu=False):
+ a = array_ops.ones([3, 4, 1], dtype=dtypes.int32)
+ b = sparse_tensor.SparseTensor([[0, 0, 1, 0], [0, 0, 3, 0]], [10, 20],
+ [1, 1, 4, 2])
+ c = a * b
+ with self.assertRaisesRegexp(
+ errors.InvalidArgumentError,
+ "broadcasts dense to sparse only; got incompatible shapes"):
+ c.eval()
+
+ @test_util.run_deprecated_v1
def testCwiseDivAndMul(self):
np.random.seed(1618)
sp_shapes = [(10, 10, 10), (5, 5), (1618,), (3, 3, 7)]
diff --git a/tensorflow/python/kernel_tests/stack_ops_test.py b/tensorflow/python/kernel_tests/stack_ops_test.py
index d50f3f4..1930d24 100644
--- a/tensorflow/python/kernel_tests/stack_ops_test.py
+++ b/tensorflow/python/kernel_tests/stack_ops_test.py
@@ -96,7 +96,7 @@
c1, b1, [r, v], [r.get_shape(), tensor_shape.unknown_shape()])
self.assertAllClose(np.ones(2000) * 10.0, self.evaluate(ry))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testStackWhileSwap(self):
self._testStackWhileSwap(use_gpu=False)
self._testStackWhileSwap(use_gpu=True)
@@ -248,7 +248,7 @@
c1, b1, [r, v], [r.get_shape(), tensor_shape.unknown_shape()])
self.assertAllClose(np.ones(2000) * 10.0, self.evaluate(ry))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testStackWhileSwap(self):
self._testStackWhileSwap(use_gpu=False)
self._testStackWhileSwap(use_gpu=True)
diff --git a/tensorflow/python/kernel_tests/svd_op_test.py b/tensorflow/python/kernel_tests/svd_op_test.py
index 97a280e..cfa9f12 100644
--- a/tensorflow/python/kernel_tests/svd_op_test.py
+++ b/tensorflow/python/kernel_tests/svd_op_test.py
@@ -20,7 +20,9 @@
import numpy as np
+from tensorflow.python import tf2
from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradient_checker
from tensorflow.python.ops import linalg_ops
@@ -38,6 +40,7 @@
class SvdOpTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testWrongDimensions(self):
# The input to svd should be a tensor of at least rank 2.
scalar = constant_op.constant(1.)
@@ -49,6 +52,7 @@
"Shape must be at least rank 2 but is rank 1"):
linalg_ops.svd(vector)
+ @test_util.run_v1_only("b/120545219")
def testConcurrentExecutesWithoutError(self):
with self.session(use_gpu=True) as sess:
all_ops = []
@@ -117,14 +121,15 @@
diag_s = array_ops.concat([diag_s, zeros], a.ndim - 1)
a_recon = math_ops.matmul(u, diag_s)
a_recon = math_ops.matmul(a_recon, v, adjoint_b=True)
- self.assertAllClose(a_recon.eval(), a, rtol=tol, atol=tol)
+ self.assertAllClose(a_recon, a, rtol=tol, atol=tol)
def CheckUnitary(self, x, tol):
# Tests that x[...,:,:]^H * x[...,:,:] is close to the identity.
xx = math_ops.matmul(x, x, adjoint_a=True)
identity = array_ops.matrix_band_part(array_ops.ones_like(xx), 0, 0)
- self.assertAllClose(identity.eval(), self.evaluate(xx), atol=tol)
+ self.assertAllClose(identity, xx, atol=tol)
+ @test_util.run_v1_only("b/120545219")
def Test(self):
is_complex = dtype_ in (np.complex64, np.complex128)
is_single = dtype_ in (np.float32, np.complex64)
@@ -213,6 +218,7 @@
tf_v *= phase[..., :n]
return tf_s, tf_u, tf_v
+ @test_util.run_v1_only("b/120545219")
def Test(self):
np.random.seed(42)
a = np.random.uniform(low=-1.0, high=1.0, size=shape_).astype(dtype_)
@@ -263,7 +269,8 @@
for cols in 1, 2, 5, 10, 32, 100:
for batch_dims in [(), (3,)] + [(3, 2)] * (max(rows, cols) < 10):
shape = batch_dims + (rows, cols)
- for use_static_shape in True, False:
+ # TF2 does not support placeholders under eager so we skip it
+ for use_static_shape in set([True, tf2.enabled()]):
name = "%s_%s_static_shape_%s__compute_uv_%s_full_%s" % (
dtype.__name__, "_".join(map(str, shape)), use_static_shape,
compute_uv, full_matrices)
diff --git a/tensorflow/python/kernel_tests/tensor_array_ops_test.py b/tensorflow/python/kernel_tests/tensor_array_ops_test.py
index 884c04e..147e7fd 100644
--- a/tensorflow/python/kernel_tests/tensor_array_ops_test.py
+++ b/tensorflow/python/kernel_tests/tensor_array_ops_test.py
@@ -32,6 +32,7 @@
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 control_flow_util
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import gen_data_flow_ops
from tensorflow.python.ops import gradients_impl
@@ -309,7 +310,7 @@
self._testTensorArraySplitRead(dtypes.string)
@test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("v2 does not support TensorArray.grad.")
def testSkipEagerTensorGradArrayWriteRead(self):
with self.session(use_gpu=True) as session:
ta = tensor_array_ops.TensorArray(
@@ -345,7 +346,7 @@
@test_util.run_deprecated_v1
def testSkipEagerTensorArrayGradGrad(self):
- if not tensor_array_ops.ENABLE_TENSOR_ARRAY_V2:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
self.skipTest("Legacy TensorArray does not support double derivatives.")
with self.test_session(use_gpu=True) as session:
x = constant_op.constant(4.0)
@@ -364,7 +365,7 @@
self.assertAllEqual([2.0], session.run(g2))
@test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("v2 does not support TensorArray.grad.")
def testSkipEagerTensorGradArrayDynamicWriteRead(self):
with self.session(use_gpu=True) as session:
ta = tensor_array_ops.TensorArray(
@@ -407,7 +408,7 @@
self.assertAllEqual(3, g_vs)
@test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("v2 does not support TensorArray.grad.")
def testSkipEagerTensorGradAccessTwiceReceiveSameObject(self):
with self.session(use_gpu=True) as session:
ta = tensor_array_ops.TensorArray(
@@ -424,12 +425,11 @@
self.assertAllEqual(t_g_ta_0, t_g_ta_1)
self.assertAllEqual([[4.0, 5.0]], d_r1_0)
- @test_util.run_deprecated_v1
def testTensorArrayWriteWrongIndexOrDataTypeFails(self):
with self.session(use_gpu=True):
ta = _make_ta(3, "foo", dtype=dtypes.float32)
# Test writing the wrong datatype
- if (tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and
+ if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
not context.executing_eagerly()):
error_msg = ("Invalid data types; op elements string but list elements "
"float")
@@ -440,7 +440,7 @@
with self.assertRaisesOpError(error_msg):
self.evaluate(ta.write(0, "wrong_type_scalar").flow)
- if (tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and
+ if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
not context.executing_eagerly()):
error_msg = "Trying to modify element -1 in a list with 3 elements."
else:
@@ -448,7 +448,7 @@
with self.assertRaisesOpError(error_msg):
self.evaluate(ta.write(-1, 3.0).flow)
- if (tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and
+ if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
not context.executing_eagerly()):
error_msg = "Trying to modify element 3 in a list with 3 elements"
else:
@@ -458,7 +458,6 @@
with self.assertRaisesOpError(error_msg):
self.evaluate(ta.write(3, 3.0).flow)
- @test_util.run_deprecated_v1
def testTensorArrayReadWrongIndexOrDataTypeFails(self):
with self.session(use_gpu=True):
ta = _make_ta(3, "foo", dtype=dtypes.float32)
@@ -467,14 +466,14 @@
# Test reading wrong datatype (only possible when constructing graphs).
if (not context.executing_eagerly() and
- not tensor_array_ops.ENABLE_TENSOR_ARRAY_V2):
+ not control_flow_util.ENABLE_CONTROL_FLOW_V2):
r0_bad = gen_data_flow_ops.tensor_array_read_v3(
handle=w0.handle, index=0, dtype=dtypes.float64, flow_in=w0.flow)
with self.assertRaisesOpError(
"TensorArray dtype is float but Op requested dtype double."):
self.evaluate(r0_bad)
- if (tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and
+ if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
not context.executing_eagerly()):
error_msg = "Trying to access element -1 in a list with 3 elements."
else:
@@ -483,7 +482,7 @@
with self.assertRaisesOpError(error_msg):
self.evaluate(ta.read(-1))
- if (tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and
+ if (control_flow_util.ENABLE_CONTROL_FLOW_V2 and
not context.executing_eagerly()):
error_msg = "Trying to access element 3 in a list with 3 elements."
else:
@@ -493,7 +492,7 @@
self.evaluate(ta.read(3))
@test_util.disable_control_flow_v2("v2 allows multiple writes.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("v2 allows multiple writes.")
def testSkipEagerTensorArrayWriteMultipleFails(self):
with self.session(use_gpu=True):
ta = tensor_array_ops.TensorArray(
@@ -504,7 +503,6 @@
"it has already been written to."):
self.evaluate(ta.write(2, 3.0).write(2, 3.0).flow)
- @test_util.run_deprecated_v1
def testTensorArrayConcatIncompatibleShapesFails(self):
with self.session(use_gpu=True):
ta = tensor_array_ops.TensorArray(
@@ -536,7 +534,6 @@
with self.assertRaisesOpError("shape"):
self.evaluate(w3.concat())
- @test_util.run_deprecated_v1
def testTensorArraySplitIncompatibleShapesFails(self):
with self.session(use_gpu=True):
in_eager_mode = context.executing_eagerly()
@@ -550,7 +547,7 @@
ta.split([1.0, 2.0, 3.0], lengths).flow.eval(feed_dict={lengths: 1})
error_msg = ("Unused values in tensor. Length of tensor: 3 Values used: 1"
- if tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2 and
not in_eager_mode else
r"Expected sum of lengths to be equal to values.shape\[0\], "
r"but sum of lengths is 1 and value's shape is: \[3\]")
@@ -558,7 +555,7 @@
self.evaluate(ta.split([1.0, 2.0, 3.0], [1]).flow)
ta = _make_ta(1, "baz")
- if tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 and not in_eager_mode:
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2 and not in_eager_mode:
with self.assertRaisesRegexp(
ValueError, "Shape must be at least rank 1 but is rank 0"):
self.evaluate(ta.split(1.0, [1]).flow)
@@ -568,7 +565,7 @@
):
self.evaluate(ta.split(1.0, [1]).flow)
- if not tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 or in_eager_mode:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2 or in_eager_mode:
ta = _make_ta(2, "buz")
with self.assertRaisesOpError(
r"TensorArray's size is not equal to the size of lengths "
@@ -611,14 +608,14 @@
wb1_grad.flow.eval()
@test_util.disable_control_flow_v2("v2 does not support TensorArray.grad.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("v2 does not support TensorArray.grad.")
def testSkipEagerTensorArrayWriteGradientAddMultipleAdds(self):
for dtype in (dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64,
dtypes.complex64, dtypes.complex128):
self._testTensorArrayWriteGradientAddMultipleAdds(dtype)
@test_util.disable_control_flow_v2("Low level legacy TA op test.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("Low level legacy TA op test.")
def testSkipEagerTensorArrayGradWithShapeKnownElementShape(self):
with self.session(use_gpu=True) as sess:
ta = tensor_array_ops.TensorArray(
@@ -649,7 +646,7 @@
sess.run(read_value, feed_dict={value: fed_value}))
@test_util.disable_control_flow_v2("Low level legacy TA op test.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("Low level legacy TA op test.")
def testSkipEagerTensorArrayGradWithShapeUnknownElementShape(self):
with self.session(use_gpu=True) as sess:
ta = tensor_array_ops.TensorArray(
@@ -778,7 +775,7 @@
self._testTensorArrayGradientWritePackConcatAndRead()
@test_util.disable_control_flow_v2("v2 does not support clear_after_read.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("v2 does not support clear_after_read.")
def testTensorArrayReadTwice(self):
with self.session(use_gpu=True):
value = constant_op.constant([[1.0, -1.0], [10.0, -10.0]])
@@ -958,7 +955,7 @@
v0_grad = gradients_impl.gradients([vout], [v0], [grad_val])[0]
state0_grad = gradients_impl.gradients([vout], [state0], [grad_val])[0]
var_grad = gradients_impl.gradients([vout], [var], [grad_val])[0]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
state0_t, var_t, v0_t, vout_t, v0_grad_t, var_grad_t, state0_grad_t = (
self.evaluate(
@@ -1003,21 +1000,6 @@
# self._testWhileLoopWritePackGradients(
# dynamic_size=False, dtype=tf.int64)
- @test_util.disable_control_flow_v2("Testing v1 while_loop with v2 TA")
- @test_util.enable_tensor_array_v2
- def testWhileLoopV1WithTensorArrayV2(self):
- size = 3
- ta = tensor_array_ops.TensorArray(
- dtype=dtypes.int32, size=size, element_shape=tensor_shape.scalar())
-
- def Body(counter, ta):
- return counter + 1, ta.write(counter, counter)
-
- _, ta = control_flow_ops.while_loop(lambda i, _: i < size, Body, [0, ta])
-
- for i in range(size):
- self.assertEqual(self.evaluate(ta.read(i)), i)
-
@test_util.disable_control_flow_v2("b/117943489 (dynamic_size)")
@test_util.run_v1_only("b/117943489")
def testSkipEagerWhileLoopDynamicWritePackGradients(self):
@@ -1270,7 +1252,7 @@
self.assertEqual((2, 2), w0.read(1).get_shape())
else:
self.assertEqual(r0.get_shape().ndims, None)
- if not tensor_array_ops.ENABLE_TENSOR_ARRAY_V2:
+ if not control_flow_util.ENABLE_CONTROL_FLOW_V2:
self.assertEqual(
tensor_shape.TensorShape(
ta1.handle.op.get_attr("element_shape")).ndims, None)
@@ -1347,11 +1329,11 @@
"TensorArray has size zero, but element shape <unknown> is not "
"fully defined. Currently only static shapes are supported when "
"packing zero-size TensorArrays.")
- with self.assertRaisesOpError(v2_msg if tensor_array_ops
- .ENABLE_TENSOR_ARRAY_V2 else v1_msg):
+ with self.assertRaisesOpError(
+ v2_msg if control_flow_util.ENABLE_CONTROL_FLOW_V2 else v1_msg):
ta.stack().eval()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSkipEagerTensorArrayEvalEmpty(self):
self._testTensorArrayEvalEmpty()
@@ -1443,7 +1425,7 @@
self.assertAllEqual(expected_grad, grad_vals[0])
@test_util.disable_control_flow_v2("colocate_with not supported in v2.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSkipEagerTensorArrayGetsDeviceFromFirstWrite(self):
with ops.device("/job:worker/task:0/cpu:0"):
# this initial device will be ignored.
@@ -1493,7 +1475,7 @@
[s for s in dev_stats[d] if "/TensorArray" in s.node_name])
@test_util.disable_control_flow_v2("colocate_with not supported in v2.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSkipEagerTensorArrayGetsDeviceFromFirstWriteInWhileLoop(self):
with ops.device("/job:worker/task:0/cpu:0"):
ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2)
@@ -1524,7 +1506,7 @@
[s for s in dev_stats[d] if "TensorArray" == s.node_name])
@test_util.disable_control_flow_v2("colocate_with not supported in v2.")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSkipEagerTensorArrayDisabledColocateWithFirstWriteCall(self):
with ops.device("/job:worker/task:0/cpu:0"):
ta = tensor_array_ops.TensorArray(
@@ -1592,7 +1574,7 @@
self.assertEqual(tensor_shape.scalar(), read1.get_shape())
if not context.executing_eagerly():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
read0_v, read1_v, size0_v, size1_v = self.evaluate((read0, read1, size0,
size1))
diff --git a/tensorflow/python/kernel_tests/tensordot_op_test.py b/tensorflow/python/kernel_tests/tensordot_op_test.py
index 123c9b3..febfe23 100644
--- a/tensorflow/python/kernel_tests/tensordot_op_test.py
+++ b/tensorflow/python/kernel_tests/tensordot_op_test.py
@@ -20,9 +20,11 @@
import numpy as np
+from tensorflow.python import tf2
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 test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test as test_lib
@@ -39,6 +41,7 @@
class TensordotTest(test_lib.TestCase):
+ @test_util.run_v1_only("b/120545219")
def test_invalid_shape(self):
a = [[1, 2], [3, 4]]
b = [[1, 2], [3, 4], [5, 6]]
@@ -62,6 +65,7 @@
axes_ph: (a_axes, b_axes)
})
+ @test_util.run_v1_only("b/120545219")
def test_invalid_axes(self):
a = [[1, 2], [3, 4]]
b = [[1, 2], [3, 4]]
@@ -99,11 +103,12 @@
tf_a = array_ops.ones((3, 3), dtype=dtypes.float32)
tf_b = constant_op.constant([2, 3, 1], dtype=dtypes.float32)[None, None]
- tf_ans = math_ops.tensordot(tf_a, tf_b, axes_value).eval()
+ tf_ans = math_ops.tensordot(tf_a, tf_b, axes_value)
self.assertAllEqual(tf_ans.shape, np_ans.shape)
self.assertAllEqual(tf_ans, np_ans)
+ @test_util.run_v1_only("b/120545219")
def test_partial_shape_inference(self):
for axes in ([1], [0]), 1:
a = array_ops.placeholder(dtypes.float32)
@@ -178,7 +183,7 @@
axes: (a_dims_np, b_dims_np)
})
else:
- tf_ans = math_ops.tensordot(a_np, b_np, (a_dims_np, b_dims_np)).eval()
+ tf_ans = math_ops.tensordot(a_np, b_np, (a_dims_np, b_dims_np))
self.assertAllClose(tf_ans, np_ans, rtol=tol, atol=tol)
self.assertAllEqual(tf_ans.shape, np_ans.shape)
@@ -208,7 +213,7 @@
c = math_ops.tensordot(a, b, axes=axes)
tf_ans = sess.run(c, feed_dict={a: a_np, b: b_np})
else:
- tf_ans = math_ops.tensordot(a_np, b_np, axes=axes).eval()
+ tf_ans = math_ops.tensordot(a_np, b_np, axes=axes)
self.assertAllClose(tf_ans, np_ans, rtol=tol, atol=tol)
self.assertAllEqual(tf_ans.shape, np_ans.shape)
@@ -220,7 +225,8 @@
for rank_a in 1, 2, 4, 5:
for rank_b in 1, 2, 4, 5:
for num_dims in range(0, min(rank_a, rank_b) + 1):
- for dynamic_shape in False, True:
+ # TF2 does not support placeholders under eager so we skip it
+ for dynamic_shape in set([False, not tf2.enabled()]):
for testcase in _get_tensordot_tests(dtype, rank_a, rank_b,
num_dims, dynamic_shape):
name = "%s_%s_%s_%s_%s_%s" % (testcase.__name__, dtype.__name__,
diff --git a/tensorflow/python/kernel_tests/unicode_decode_op_test.py b/tensorflow/python/kernel_tests/unicode_decode_op_test.py
index c165021..9a59f8a 100644
--- a/tensorflow/python/kernel_tests/unicode_decode_op_test.py
+++ b/tensorflow/python/kernel_tests/unicode_decode_op_test.py
@@ -19,134 +19,686 @@
from __future__ import division
from __future__ import print_function
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import errors_impl as errors
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_string_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_test_util
from tensorflow.python.platform import test
-# Account for python2 and python3 execution of the test.
-def codepoint(s):
- if isinstance(s, bytes):
- return ord(s.decode("utf-8"))
- elif isinstance(s, str):
- return ord(s)
+def _nested_encode(x, encoding):
+ """Encode each string in a nested list with `encoding`."""
+ if isinstance(x, list):
+ return [_nested_encode(v, encoding) for v in x]
+ else:
+ return x.encode(encoding)
-class UnicodeDecodeTest(test.TestCase):
+def _nested_codepoints(x):
+ """Replace each string in a nested list with a list of its codepoints."""
+ # Works for Python 2 and 3, and for both UCS2 and UCS4 builds
+ if isinstance(x, list):
+ return [_nested_codepoints(v) for v in x]
+ else:
+ b = list(x.encode("utf-32-be"))
+ if any(isinstance(c, str) for c in b):
+ b = [ord(c) for c in b]
+ return [(b0 << 24) + (b1 << 16) + (b2 << 8) + b3
+ for b0, b1, b2, b3 in zip(b[::4], b[1::4], b[2::4], b[3::4])]
- def testBatchDecode(self):
- text = constant_op.constant(
- ["仅今年前", "分享介面終於迎來更新"])
- row_splits, utf8_text, offsets = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8")
- with self.test_session():
- self.assertAllEqual([
- codepoint("仅"),
- codepoint("今"),
- codepoint("年"),
- codepoint("前"),
- codepoint("分"),
- codepoint("享"),
- codepoint("介"),
- codepoint("面"),
- codepoint("終"),
- codepoint("於"),
- codepoint("迎"),
- codepoint("來"),
- codepoint("更"),
- codepoint("新")
- ],
- self.evaluate(utf8_text).tolist())
- self.assertAllEqual([0, 4, 14], self.evaluate(row_splits).tolist())
- self.assertAllEqual([0, 3, 6, 9, 0, 3, 6, 9, 12, 15, 18, 21, 24, 27],
- self.evaluate(offsets).tolist())
+def _nested_offsets(x, encoding):
+ """Replace each string in a nested list with a list of start offsets."""
+ if isinstance(x, list):
+ return [_nested_offsets(v, encoding) for v in x]
+ else:
+ if not x:
+ return []
+ encoded_x = x.encode("utf-32-be")
+ encoded_chars = [encoded_x[i:i + 4] for i in range(0, len(encoded_x), 4)]
+ char_lens = [
+ len(c.decode("utf-32-be").encode(encoding)) for c in encoded_chars
+ ]
+ return [0] + np.cumsum(char_lens).tolist()[:-1]
- def testBasicDecodeWithOffset(self):
- text = constant_op.constant(["仅今年前"])
- row_splits, utf8_text, starts = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8")
- with self.test_session():
- self.assertAllEqual([
- codepoint("仅"),
- codepoint("今"),
- codepoint("年"),
- codepoint("前"),
- ],
- self.evaluate(utf8_text).tolist())
- self.assertAllEqual(self.evaluate(row_splits).tolist(), [0, 4])
- self.assertAllEqual(self.evaluate(starts).tolist(), [0, 3, 6, 9])
+def _nested_splitchars(x, encoding):
+ """Replace each string in a nested list with a list of char substrings."""
+ if isinstance(x, list):
+ return [_nested_splitchars(v, encoding) for v in x]
+ else:
+ b = x.encode("utf-32-be")
+ chars = zip(b[::4], b[1::4], b[2::4], b[3::4])
+ if str is bytes:
+ return [b"".join(c).decode("utf-32-be").encode(encoding) for c in chars]
+ else:
+ return [bytes(c).decode("utf-32-be").encode(encoding) for c in chars]
- @test_util.run_deprecated_v1
- def testStrictError(self):
- text = constant_op.constant([b"\xFEED"])
- _, error, _ = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8", errors="strict")
- with self.assertRaises(errors.InvalidArgumentError):
- with self.test_session():
- self.evaluate(error)
+def _make_sparse_tensor(indices, values, dense_shape, dtype=np.int32):
+ return sparse_tensor.SparseTensorValue(
+ np.array(indices, np.int64), np.array(values, dtype),
+ np.array(dense_shape, np.int64))
- def testReplaceOnError(self):
- text = constant_op.constant([b"\xFE"])
- _, utf8_text, _ = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8", errors="replace")
+@test_util.run_all_in_graph_and_eager_modes
+class UnicodeDecodeTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
- with self.test_session():
- self.assertAllEqual(self.evaluate(utf8_text).tolist(), [65533])
+ def testScalarDecode(self):
+ text = constant_op.constant(u"仅今年前".encode("utf-8"))
+ chars = ragged_string_ops.unicode_decode(text, "utf-8")
+ self.assertAllEqual(chars, [ord(c) for c in u"仅今年前"])
- @test_util.run_deprecated_v1
- def testBadReplacementChar(self):
- text = constant_op.constant([b"\xFE"])
- _, error, _ = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8", errors="replace", replacement_char=11141111)
+ def testScalarDecodeWithOffset(self):
+ text = constant_op.constant(u"仅今年前".encode("utf-8"))
+ chars, starts = ragged_string_ops.unicode_decode_with_offsets(text, "utf-8")
+ self.assertAllEqual(chars, [ord(c) for c in u"仅今年前"])
+ self.assertAllEqual(starts, [0, 3, 6, 9])
- with self.assertRaises(errors.InvalidArgumentError):
- with self.test_session():
- self.evaluate(error)
+ def testVectorDecode(self):
+ text = constant_op.constant([u"仅今年前".encode("utf-8"), b"hello"])
+ chars = ragged_string_ops.unicode_decode(text, "utf-8")
+ expected_chars = [[ord(c) for c in u"仅今年前"],
+ [ord(c) for c in u"hello"]]
+ self.assertRaggedEqual(chars, expected_chars)
- def testIgnoreOnError(self):
- text = constant_op.constant([b"\xFEhello"])
+ def testVectorDecodeWithOffset(self):
+ text = constant_op.constant([u"仅今年前".encode("utf-8"), b"hello"])
+ chars, starts = ragged_string_ops.unicode_decode_with_offsets(text, "utf-8")
+ expected_chars = [[ord(c) for c in u"仅今年前"],
+ [ord(c) for c in u"hello"]]
+ self.assertRaggedEqual(chars, expected_chars)
+ self.assertRaggedEqual(starts, [[0, 3, 6, 9], [0, 1, 2, 3, 4]])
- _, utf8_text, _ = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8", errors="ignore")
+ @parameterized.parameters([
+ {"texts": u"仅今年前"},
+ {"texts": [u"G\xf6\xf6dnight", u"\U0001f60a"]},
+ {"texts": ["Hello", "world", "", u"👍"]},
+ {"texts": [["Hi", "there"], ["", u"\U0001f60a"]], "ragged_rank": 0},
+ {"texts": [["Hi", "there", ""], [u"😊"]], "ragged_rank": 1},
+ {"texts": [[[u"😊", u"🤠🧐"], []], [[u"🤓👻🤖"]]], "ragged_rank": 2},
+ {"texts": []}
+ ]) # pyformat: disable
+ def testBasicDecode(self, texts, ragged_rank=None):
+ input_tensor = ragged_factory_ops.constant_value(
+ _nested_encode(texts, "UTF-8"), ragged_rank=ragged_rank, dtype=bytes)
+ result = ragged_string_ops.unicode_decode(input_tensor, "UTF-8")
+ expected = _nested_codepoints(texts)
+ self.assertRaggedEqual(expected, result)
- with self.test_session():
- self.assertAllEqual(self.evaluate(utf8_text).tolist(), [
- codepoint("h"),
- codepoint("e"),
- codepoint("l"),
- codepoint("l"),
- codepoint("o")
- ])
+ @parameterized.parameters([
+ {"texts": u"仅今年前"},
+ {"texts": [u"G\xf6\xf6dnight", u"\U0001f60a"]},
+ {"texts": ["Hello", "world", "", u"👍"]},
+ {"texts": [["Hi", "there"], ["", u"\U0001f60a"]], "ragged_rank": 0},
+ {"texts": [["Hi", "there", ""], [u"😊"]], "ragged_rank": 1},
+ {"texts": [[[u"😊", u"🤠🧐"], []], [[u"🤓👻🤖"]]], "ragged_rank": 2},
+ {"texts": []}
+ ]) # pyformat: disable
+ def testBasicDecodeWithOffsets(self, texts, ragged_rank=None):
+ input_tensor = ragged_factory_ops.constant_value(
+ _nested_encode(texts, "UTF-8"), ragged_rank=ragged_rank, dtype=bytes)
+ result = ragged_string_ops.unicode_decode_with_offsets(
+ input_tensor, "UTF-8")
+ expected_codepoints = _nested_codepoints(texts)
+ expected_offsets = _nested_offsets(texts, "UTF-8")
+ self.assertRaggedEqual(expected_codepoints, result[0])
+ self.assertRaggedEqual(expected_offsets, result[1])
- @test_util.run_deprecated_v1
- def testBadErrorPolicy(self):
- text = constant_op.constant(["hippopotamus"])
+ def testDocstringExamples(self):
+ texts = [s.encode("utf8") for s in [u"G\xf6\xf6dnight", u"\U0001f60a"]]
+ codepoints1 = ragged_string_ops.unicode_decode(texts, "UTF-8")
+ codepoints2, offsets = ragged_string_ops.unicode_decode_with_offsets(
+ texts, "UTF-8")
+ self.assertRaggedEqual(
+ codepoints1, [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]])
+ self.assertRaggedEqual(
+ codepoints2, [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]])
+ self.assertRaggedEqual(offsets, [[0, 1, 3, 5, 6, 7, 8, 9, 10], [0]])
- with self.assertRaises(ValueError):
- _, _, _ = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8", errors="oranguatan")
+ @parameterized.parameters([
+ dict(
+ texts=["Hello", "world", "", u"👍"],
+ expected=_make_sparse_tensor(
+ indices=[[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [1, 0], [1, 1],
+ [1, 2], [1, 3], [1, 4], [3, 0]],
+ values=[72, 101, 108, 108, 111, 119, 111, 114, 108, 100, 128077],
+ dense_shape=[4, 5])),
+ dict(
+ texts=[["Hi", "there"], ["", u"\U0001f60a"]],
+ expected=_make_sparse_tensor(
+ indices=[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [0, 1, 2],
+ [0, 1, 3], [0, 1, 4], [1, 1, 0]],
+ values=[72, 105, 116, 104, 101, 114, 101, 128522],
+ dense_shape=[2, 2, 5])),
+ dict(
+ texts=[],
+ expected=_make_sparse_tensor(np.zeros([0, 2], np.int64), [], [0, 0])),
+ ])
+ def testDecodeWithSparseOutput(self, texts, expected):
+ input_tensor = np.array(_nested_encode(texts, "UTF-8"), dtype=bytes)
+ result = ragged_string_ops.unicode_decode(input_tensor, "UTF-8").to_sparse()
+ self.assertIsInstance(result, sparse_tensor.SparseTensor)
+ self.assertAllEqual(expected.indices, result.indices)
+ self.assertAllEqual(expected.values, result.values)
+ self.assertAllEqual(expected.dense_shape, result.dense_shape)
- def testReplaceControlChars(self):
- text = constant_op.constant(["\x02仅今年前"])
- row_splits, utf8_text, _ = gen_string_ops.unicode_decode_with_offsets(
- text, "utf-8", replace_control_characters=True)
+ @parameterized.parameters([
+ dict(
+ texts=["Hello", "world", "", u"👍"],
+ expected=[[72, 101, 108, 108, 111], [119, 111, 114, 108, 100],
+ [-1, -1, -1, -1, -1], [128077, -1, -1, -1, -1]]),
+ dict(
+ texts=[["Hi", "there"], ["", u"\U0001f60a"]],
+ expected=[[[72, 105, -1, -1, -1], [116, 104, 101, 114, 101]],
+ [[-1, -1, -1, -1, -1], [128522, -1, -1, -1, -1]]],
+ ragged_rank=0),
+ dict(
+ texts=[["Hi", "there", ""], [u"😊"]],
+ expected=[[[72, 105, -1, -1, -1],
+ [116, 104, 101, 114, 101],
+ [-1, -1, -1, -1, -1]],
+ [[128522, -1, -1, -1, -1],
+ [-1, -1, -1, -1, -1],
+ [-1, -1, -1, -1, -1]]]),
+ dict(
+ texts=[[[u"😊", u"🤠🧐"], []], [[u"🤓👻🤖"]]],
+ expected=[
+ [[[128522, -1, -1], [129312, 129488, -1]],
+ [[-1, -1, -1], [-1, -1, -1]]],
+ [[[129299, 128123, 129302], [-1, -1, -1]],
+ [[-1, -1, -1], [-1, -1, -1]]]]),
+ dict(texts=[], expected=np.zeros([0, 0], np.int64)),
+ ]) # pyformat: disable
+ def testDecodeWithPaddedOutput(self, texts, expected, ragged_rank=None):
+ input_tensor = ragged_factory_ops.constant_value(
+ _nested_encode(texts, "UTF-8"), ragged_rank=ragged_rank, dtype=bytes)
+ result = ragged_string_ops.unicode_decode(
+ input_tensor, "UTF-8").to_tensor(default_value=-1)
+ self.assertAllEqual(expected, result)
- with self.test_session():
- self.assertAllEqual([
- 65533,
- codepoint("仅"),
- codepoint("今"),
- codepoint("年"),
- codepoint("前"),
- ],
- self.evaluate(utf8_text).tolist())
- self.assertAllEqual([0, 5], self.evaluate(row_splits).tolist())
+ @parameterized.parameters([
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ expected=[[65533], [104, 101, 108, 108, 111],
+ [61, 61, 65533, 61, 61], [119, 111, 114, 108, 100]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ replacement_char=0,
+ expected=[[0], [104, 101, 108, 108, 111],
+ [61, 61, 0, 61, 61], [119, 111, 114, 108, 100]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="ignore",
+ expected=[[], [104, 101, 108, 108, 111],
+ [61, 61, 61, 61], [119, 111, 114, 108, 100]]),
+ dict(
+ input=[b"\x00", b"hello", b"==\x01==", b"world"],
+ input_encoding="UTF-8",
+ replace_control_characters=True,
+ expected=[[65533], [104, 101, 108, 108, 111],
+ [61, 61, 65533, 61, 61], [119, 111, 114, 108, 100]]),
+ dict(
+ input=[b"\x00", b"hello", b"==\x01==", b"world"],
+ input_encoding="UTF-8",
+ replace_control_characters=True,
+ replacement_char=0,
+ expected=[[0], [104, 101, 108, 108, 111],
+ [61, 61, 0, 61, 61], [119, 111, 114, 108, 100]]),
+ ]) # pyformat: disable
+ def testErrorModes(self, expected=None, **args):
+ result = ragged_string_ops.unicode_decode(**args)
+ self.assertRaggedEqual(expected, result)
+
+ @parameterized.parameters([
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ expected=[[65533], [104, 101, 108, 108, 111],
+ [61, 61, 65533, 61, 61], [119, 111, 114, 108, 100]],
+ expected_offsets=[[0], [0, 1, 2, 3, 4],
+ [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ replacement_char=0,
+ expected=[[0], [104, 101, 108, 108, 111],
+ [61, 61, 0, 61, 61], [119, 111, 114, 108, 100]],
+ expected_offsets=[[0], [0, 1, 2, 3, 4],
+ [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="ignore",
+ expected=[[], [104, 101, 108, 108, 111],
+ [61, 61, 61, 61], [119, 111, 114, 108, 100]],
+ expected_offsets=[[], [0, 1, 2, 3, 4],
+ [0, 1, 3, 4], [0, 1, 2, 3, 4]]),
+ dict(
+ input=[b"\x00", b"hello", b"==\x01==", b"world"],
+ input_encoding="UTF-8",
+ replace_control_characters=True,
+ expected=[[65533], [104, 101, 108, 108, 111],
+ [61, 61, 65533, 61, 61], [119, 111, 114, 108, 100]],
+ expected_offsets=[[0], [0, 1, 2, 3, 4],
+ [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]),
+ dict(
+ input=[b"\x00", b"hello", b"==\x01==", b"world"],
+ input_encoding="UTF-8",
+ replace_control_characters=True,
+ replacement_char=0,
+ expected=[[0], [104, 101, 108, 108, 111],
+ [61, 61, 0, 61, 61], [119, 111, 114, 108, 100]],
+ expected_offsets=[[0], [0, 1, 2, 3, 4],
+ [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]),
+ ]) # pyformat: disable
+ def testErrorModesWithOffsets(self,
+ expected=None,
+ expected_offsets=None,
+ **args):
+ result = ragged_string_ops.unicode_decode_with_offsets(**args)
+ self.assertRaggedEqual(result[0], expected)
+ self.assertRaggedEqual(result[1], expected_offsets)
+
+ @parameterized.parameters(
+ ("UTF-8", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-16-BE", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-32-BE", [u"こんにちは", u"你好", u"Hello"]),
+ ("US-ASCII", [u"Hello", "world"]),
+ ("ISO-8859-1", [u"ÀÈÓ", "AEO"]),
+ ("SHIFT-JIS", [u"Hello", u"こんにちは"]),
+ )
+ def testDecodeWithDifferentEncodings(self, encoding, texts):
+ expected = _nested_codepoints(texts)
+ input_tensor = constant_op.constant(_nested_encode(texts, encoding))
+ result = ragged_string_ops.unicode_decode(input_tensor, encoding)
+ self.assertRaggedEqual(expected, result)
+
+ @parameterized.parameters(
+ ("UTF-8", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-16-BE", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-32-BE", [u"こんにちは", u"你好", u"Hello"]),
+ ("US-ASCII", [u"Hello", "world"]),
+ ("ISO-8859-1", [u"ÀÈÓ", "AEO"]),
+ ("SHIFT-JIS", [u"Hello", u"こんにちは"]),
+ )
+ def testDecodeWithOffsetsWithDifferentEncodings(self, encoding, texts):
+ expected_codepoints = _nested_codepoints(texts)
+ expected_offsets = _nested_offsets(texts, encoding)
+ input_tensor = constant_op.constant(_nested_encode(texts, encoding))
+ result = ragged_string_ops.unicode_decode_with_offsets(
+ input_tensor, encoding)
+ self.assertRaggedEqual(expected_codepoints, result[0])
+ self.assertRaggedEqual(expected_offsets, result[1])
+
+ @parameterized.parameters([
+ dict(input=[b"\xFEED"],
+ errors="strict",
+ input_encoding="UTF-8",
+ exception=errors.InvalidArgumentError,
+ message="Invalid formatting on input string"),
+ dict(input="x",
+ input_encoding="UTF-8",
+ replacement_char=11141111,
+ exception=errors.InvalidArgumentError,
+ message="replacement_char out of unicode codepoint range"),
+ dict(input="x",
+ input_encoding="UTF-8",
+ errors="oranguatan",
+ exception=(ValueError, errors.InvalidArgumentError)),
+ ]) # pyformat: disable
+ def testExceptions(self, exception=None, message=None, **args):
+ with self.assertRaisesRegexp(exception, message):
+ self.evaluate(ragged_string_ops.unicode_decode(**args))
+
+ def testUnknownRankError(self):
+ if context.executing_eagerly():
+ return
+ s = array_ops.placeholder(dtypes.string)
+ message = "Rank of `input` must be statically known."
+ with self.assertRaisesRegexp(ValueError, message):
+ self.evaluate(ragged_string_ops.unicode_decode(s, input_encoding="UTF-8"))
+
+ @parameterized.parameters([
+ dict(
+ doc="Single string",
+ input=_nested_encode([u"仅今年前"], "utf-8"),
+ input_encoding="UTF-8",
+ expected_char_values=_nested_codepoints(u"仅今年前"),
+ expected_row_splits=[0, 4],
+ expected_char_to_byte_starts=[0, 3, 6, 9]),
+ dict(
+ doc="Multiple strings",
+ input=_nested_encode([u"仅今年前", u"你好"], "utf-8"),
+ input_encoding="UTF-8",
+ expected_char_values=_nested_codepoints(u"仅今年前你好"),
+ expected_row_splits=[0, 4, 6],
+ expected_char_to_byte_starts=[0, 3, 6, 9, 0, 3]),
+ dict(
+ doc="errors=replace",
+ input=b"=\xFE=",
+ input_encoding="UTF-8",
+ errors="replace",
+ expected_char_values=[61, 65533, 61],
+ expected_row_splits=[0, 3],
+ expected_char_to_byte_starts=[0, 1, 2]),
+ dict(
+ doc="errors=ignore",
+ input=b"=\xFE=",
+ input_encoding="UTF-8",
+ errors="ignore",
+ expected_char_values=[61, 61],
+ expected_row_splits=[0, 2],
+ expected_char_to_byte_starts=[0, 2]),
+ ])
+ def testDecodeGenOp(self,
+ doc,
+ expected_row_splits=None,
+ expected_char_values=None,
+ expected_char_to_byte_starts=None,
+ **args):
+ """Test for the c++ interface (gen_string_ops.unicode_decode)."""
+ result = gen_string_ops.unicode_decode_with_offsets(**args)
+ self.assertAllEqual(expected_row_splits, result.row_splits)
+ self.assertAllEqual(expected_char_values, result.char_values)
+ self.assertAllEqual(expected_char_to_byte_starts,
+ result.char_to_byte_starts)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class UnicodeSplitTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
+
+ def testScalarSplit(self):
+ text = constant_op.constant(u"仅今年前".encode("UTF-8"))
+ chars = ragged_string_ops.unicode_split(text, "UTF-8")
+ self.assertAllEqual(chars, [c.encode("UTF-8") for c in u"仅今年前"])
+
+ def testScalarSplitWithOffset(self):
+ text = constant_op.constant(u"仅今年前".encode("UTF-8"))
+ chars, starts = ragged_string_ops.unicode_split_with_offsets(text, "UTF-8")
+ self.assertAllEqual(chars, [c.encode("UTF-8") for c in u"仅今年前"])
+ self.assertAllEqual(starts, [0, 3, 6, 9])
+
+ def testVectorSplit(self):
+ text = constant_op.constant([u"仅今年前".encode("UTF-8"), b"hello"])
+ chars = ragged_string_ops.unicode_split(text, "UTF-8")
+ expected_chars = [[c.encode("UTF-8") for c in u"仅今年前"],
+ [c.encode("UTF-8") for c in u"hello"]]
+ self.assertRaggedEqual(chars, expected_chars)
+
+ def testVectorSplitWithOffset(self):
+ text = constant_op.constant([u"仅今年前".encode("UTF-8"), b"hello"])
+ chars, starts = ragged_string_ops.unicode_split_with_offsets(text, "UTF-8")
+ expected_chars = [[c.encode("UTF-8") for c in u"仅今年前"],
+ [c.encode("UTF-8") for c in u"hello"]]
+ self.assertRaggedEqual(chars, expected_chars)
+ self.assertRaggedEqual(starts, [[0, 3, 6, 9], [0, 1, 2, 3, 4]])
+
+ @parameterized.parameters([
+ {"texts": u"仅今年前"},
+ {"texts": [u"G\xf6\xf6dnight", u"\U0001f60a"]},
+ {"texts": ["Hello", "world", "", u"👍"]},
+ {"texts": [["Hi", "there"], ["", u"\U0001f60a"]], "ragged_rank": 0},
+ {"texts": [["Hi", "there", ""], [u"😊"]], "ragged_rank": 1},
+ {"texts": [[[u"😊", u"🤠🧐"], []], [[u"🤓👻🤖"]]], "ragged_rank": 2},
+ {"texts": []}
+ ]) # pyformat: disable
+ def testBasicSplit(self, texts, ragged_rank=None):
+ input_tensor = ragged_factory_ops.constant_value(
+ _nested_encode(texts, "UTF-8"), ragged_rank=ragged_rank, dtype=bytes)
+ result = ragged_string_ops.unicode_split(input_tensor, "UTF-8")
+ expected = _nested_splitchars(texts, "UTF-8")
+ self.assertRaggedEqual(expected, result)
+
+ @parameterized.parameters([
+ {"texts": u"仅今年前"},
+ {"texts": [u"G\xf6\xf6dnight", u"\U0001f60a"]},
+ {"texts": ["Hello", "world", "", u"👍"]},
+ {"texts": [["Hi", "there"], ["", u"\U0001f60a"]], "ragged_rank": 0},
+ {"texts": [["Hi", "there", ""], [u"😊"]], "ragged_rank": 1},
+ {"texts": [[[u"😊", u"🤠🧐"], []], [[u"🤓👻🤖"]]], "ragged_rank": 2},
+ {"texts": []}
+ ]) # pyformat: disable
+ def testBasicSplitWithOffsets(self, texts, ragged_rank=None):
+ input_tensor = ragged_factory_ops.constant_value(
+ _nested_encode(texts, "UTF-8"), ragged_rank=ragged_rank, dtype=bytes)
+ result = ragged_string_ops.unicode_split_with_offsets(input_tensor, "UTF-8")
+ expected_codepoints = _nested_splitchars(texts, "UTF-8")
+ expected_offsets = _nested_offsets(texts, "UTF-8")
+ self.assertRaggedEqual(expected_codepoints, result[0])
+ self.assertRaggedEqual(expected_offsets, result[1])
+
+ def testDocstringExamples(self):
+ texts = [s.encode("utf8") for s in [u"G\xf6\xf6dnight", u"\U0001f60a"]]
+ codepoints1 = ragged_string_ops.unicode_split(texts, "UTF-8")
+ codepoints2, offsets = ragged_string_ops.unicode_split_with_offsets(
+ texts, "UTF-8")
+ self.assertRaggedEqual(
+ codepoints1,
+ [[b"G", b"\xc3\xb6", b"\xc3\xb6", b"d", b"n", b"i", b"g", b"h", b"t"],
+ [b"\xf0\x9f\x98\x8a"]])
+ self.assertRaggedEqual(
+ codepoints2,
+ [[b"G", b"\xc3\xb6", b"\xc3\xb6", b"d", b"n", b"i", b"g", b"h", b"t"],
+ [b"\xf0\x9f\x98\x8a"]])
+ self.assertRaggedEqual(offsets, [[0, 1, 3, 5, 6, 7, 8, 9, 10], [0]])
+
+ @parameterized.parameters([
+ dict(
+ texts=["Hello", "world", "", u"👍"],
+ expected=_make_sparse_tensor(
+ indices=[[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [1, 0], [1, 1],
+ [1, 2], [1, 3], [1, 4], [3, 0]],
+ values=[b"H", b"e", b"l", b"l", b"o",
+ b"w", b"o", b"r", b"l", b"d", b"\xf0\x9f\x91\x8d"],
+ dense_shape=[4, 5],
+ dtype=bytes)),
+ dict(
+ texts=[["Hi", "there"], ["", u"\U0001f60a"]],
+ expected=_make_sparse_tensor(
+ indices=[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [0, 1, 2],
+ [0, 1, 3], [0, 1, 4], [1, 1, 0]],
+ values=[b"H", b"i", b"t", b"h", b"e", b"r", b"e",
+ b"\xf0\x9f\x98\x8a"],
+ dense_shape=[2, 2, 5],
+ dtype=bytes)),
+ dict(
+ texts=[],
+ expected=_make_sparse_tensor(
+ np.zeros([0, 2], np.int64), [], [0, 0], dtype=bytes)),
+ ]) # pyformat: disable
+ def testSplitWithSparseOutput(self, texts, expected):
+ input_tensor = np.array(_nested_encode(texts, "UTF-8"), dtype=bytes)
+ result = ragged_string_ops.unicode_split(input_tensor, "UTF-8").to_sparse()
+ self.assertIsInstance(result, sparse_tensor.SparseTensor)
+ self.assertAllEqual(expected.indices, result.indices)
+ self.assertAllEqual(expected.values, result.values)
+ self.assertAllEqual(expected.dense_shape, result.dense_shape)
+
+ @parameterized.parameters([
+ dict(
+ texts=["Hello", "world", "", u"👍"],
+ expected=[[b"H", b"e", b"l", b"l", b"o"],
+ [b"w", b"o", b"r", b"l", b"d"],
+ ["", "", "", "", ""],
+ [b"\xf0\x9f\x91\x8d", "", "", "", ""]]),
+ dict(
+ texts=[["Hi", "there"], ["", u"\U0001f60a"]],
+ expected=[[[b"H", b"i", "", "", ""],
+ [b"t", b"h", b"e", b"r", b"e"]],
+ [["", "", "", "", ""],
+ [b"\xf0\x9f\x98\x8a", "", "", "", ""]]],
+ ragged_rank=0),
+ dict(
+ texts=[["Hi", "there", ""], [u"😊"]],
+ expected=[[[b"H", b"i", "", "", ""],
+ [b"t", b"h", b"e", b"r", b"e"],
+ ["", "", "", "", ""]],
+ [[b"\xf0\x9f\x98\x8a", "", "", "", ""],
+ ["", "", "", "", ""],
+ ["", "", "", "", ""]]]),
+ dict(
+ texts=[[[u"😊", u"🤠🧐"], []], [[u"🤓👻🤖"]]],
+ expected=[[[[b"\xf0\x9f\x98\x8a", "", ""],
+ [b"\xf0\x9f\xa4\xa0", b"\xf0\x9f\xa7\x90", ""]],
+ [["", "", ""],
+ ["", "", ""]]],
+ [[[b"\xf0\x9f\xa4\x93", b"\xf0\x9f\x91\xbb",
+ b"\xf0\x9f\xa4\x96"],
+ ["", "", ""]],
+ [["", "", ""],
+ ["", "", ""]]]]),
+ dict(texts=[], expected=np.zeros([0, 0], np.int64)),
+ ]) # pyformat: disable
+ def testSplitWithPaddedOutput(self, texts, expected, ragged_rank=None):
+ input_tensor = ragged_factory_ops.constant_value(
+ _nested_encode(texts, "UTF-8"), ragged_rank=ragged_rank, dtype=bytes)
+ result = ragged_string_ops.unicode_split(
+ input_tensor, "UTF-8").to_tensor(default_value="")
+ self.assertAllEqual(np.array(expected, dtype=bytes), result)
+
+ @parameterized.parameters([
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ expected=[[b"\xef\xbf\xbd"],
+ [b"h", b"e", b"l", b"l", b"o"],
+ [b"=", b"=", b"\xef\xbf\xbd", b"=", b"="],
+ [b"w", b"o", b"r", b"l", b"d"]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ replacement_char=0,
+ expected=[[b"\x00"],
+ [b"h", b"e", b"l", b"l", b"o"],
+ [b"=", b"=", b"\x00", b"=", b"="],
+ [b"w", b"o", b"r", b"l", b"d"]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="ignore",
+ expected=[[],
+ [b"h", b"e", b"l", b"l", b"o"],
+ [b"=", b"=", b"=", b"="],
+ [b"w", b"o", b"r", b"l", b"d"]]),
+ ]) # pyformat: disable
+ def testErrorModes(self, expected=None, **args):
+ result = ragged_string_ops.unicode_split(**args)
+ self.assertRaggedEqual(expected, result)
+
+ @parameterized.parameters([
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ expected=[[b"\xef\xbf\xbd"],
+ [b"h", b"e", b"l", b"l", b"o"],
+ [b"=", b"=", b"\xef\xbf\xbd", b"=", b"="],
+ [b"w", b"o", b"r", b"l", b"d"]],
+ expected_offsets=[[0], [0, 1, 2, 3, 4],
+ [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="replace",
+ replacement_char=0,
+ expected=[[b"\x00"],
+ [b"h", b"e", b"l", b"l", b"o"],
+ [b"=", b"=", b"\x00", b"=", b"="],
+ [b"w", b"o", b"r", b"l", b"d"]],
+ expected_offsets=[[0], [0, 1, 2, 3, 4],
+ [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]),
+ dict(
+ input=[b"\xFE", b"hello", b"==\xFF==", b"world"],
+ input_encoding="UTF-8",
+ errors="ignore",
+ expected=[[],
+ [b"h", b"e", b"l", b"l", b"o"],
+ [b"=", b"=", b"=", b"="],
+ [b"w", b"o", b"r", b"l", b"d"]],
+ expected_offsets=[[], [0, 1, 2, 3, 4],
+ [0, 1, 3, 4], [0, 1, 2, 3, 4]]),
+ ]) # pyformat: disable
+ def testErrorModesWithOffsets(self,
+ expected=None,
+ expected_offsets=None,
+ **args):
+ result = ragged_string_ops.unicode_split_with_offsets(**args)
+ self.assertRaggedEqual(expected, result[0])
+ self.assertRaggedEqual(expected_offsets, result[1])
+
+ @parameterized.parameters(
+ ("UTF-8", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-16-BE", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-32-BE", [u"こんにちは", u"你好", u"Hello"]),
+ )
+ def testSplitWithDifferentEncodings(self, encoding, texts):
+ expected = _nested_splitchars(texts, encoding)
+ input_tensor = constant_op.constant(_nested_encode(texts, encoding))
+ result = ragged_string_ops.unicode_split(input_tensor, encoding)
+ self.assertRaggedEqual(expected, result)
+
+ @parameterized.parameters(
+ ("UTF-8", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-16-BE", [u"こんにちは", u"你好", u"Hello"]),
+ ("UTF-32-BE", [u"こんにちは", u"你好", u"Hello"]),
+ )
+ def testSplitWithOffsetsWithDifferentEncodings(self, encoding, texts):
+ expected_codepoints = _nested_splitchars(texts, encoding)
+ expected_offsets = _nested_offsets(texts, encoding)
+ input_tensor = constant_op.constant(_nested_encode(texts, encoding))
+ result = ragged_string_ops.unicode_split_with_offsets(
+ input_tensor, encoding)
+ self.assertRaggedEqual(expected_codepoints, result[0])
+ self.assertRaggedEqual(expected_offsets, result[1])
+
+ @parameterized.parameters([
+ dict(input=[b"\xFEED"],
+ errors="strict",
+ input_encoding="UTF-8",
+ exception=errors.InvalidArgumentError,
+ message="Invalid formatting on input string"),
+ dict(input="x",
+ input_encoding="UTF-8",
+ replacement_char=11141111,
+ exception=errors.InvalidArgumentError,
+ message="replacement_char out of unicode codepoint range"),
+ dict(input="x",
+ input_encoding="UTF-8",
+ errors="oranguatan",
+ exception=(ValueError, errors.InvalidArgumentError)),
+ ]) # pyformat: disable
+ def testExceptions(self, exception=None, message=None, **args):
+ with self.assertRaisesRegexp(exception, message):
+ self.evaluate(ragged_string_ops.unicode_split(**args))
+
+ def testUnknownRankError(self):
+ if context.executing_eagerly():
+ return
+ s = array_ops.placeholder(dtypes.string)
+ message = "Rank of `input` must be statically known."
+ with self.assertRaisesRegexp(ValueError, message):
+ self.evaluate(ragged_string_ops.unicode_decode(s, input_encoding="UTF-8"))
if __name__ == "__main__":
diff --git a/tensorflow/python/kernel_tests/unicode_encode_op_test.py b/tensorflow/python/kernel_tests/unicode_encode_op_test.py
index a5a5c20..2f3cd8a 100644
--- a/tensorflow/python/kernel_tests/unicode_encode_op_test.py
+++ b/tensorflow/python/kernel_tests/unicode_encode_op_test.py
@@ -23,14 +23,25 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors_impl as errors
-from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
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.ops.ragged import ragged_tensor_value
from tensorflow.python.platform import test
class UnicodeEncodeOpTest(test.TestCase, parameterized.TestCase):
+ def assertRaggedEqual(self, rt, expected):
+ with self.cached_session() as sess:
+ value = sess.run(rt)
+ if isinstance(value, np.ndarray):
+ value = value.tolist()
+ elif isinstance(value, ragged_tensor_value.RaggedTensorValue):
+ value = value.to_list()
+ self.assertEqual(value, expected)
+
def testScalar(self):
with self.cached_session():
with self.assertRaises(ValueError):
@@ -53,97 +64,80 @@
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
def testStrictErrors(self, encoding):
test_value = np.array([72, 101, 2147483647, -1, 111], np.int32)
- with self.cached_session():
+ with self.cached_session() as session:
with self.assertRaises(errors.InvalidArgumentError):
- ragged_string_ops.unicode_encode(test_value, encoding, "strict").eval()
+ session.run(
+ ragged_string_ops.unicode_encode(test_value, encoding, "strict"))
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def testIgnoreErrors(self, encoding):
test_value = np.array([72, 101, 2147483647, -1, 111], np.int32)
expected_value = u"Heo".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding,
"ignore")
- with self.cached_session():
- result = unicode_encode_op.eval()
+ with self.cached_session() as session:
+ result = session.run(unicode_encode_op)
self.assertIsInstance(result, bytes)
self.assertAllEqual(result, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def testReplaceErrors(self, encoding):
test_value = np.array([72, 101, 2147483647, -1, 111], np.int32)
expected_value = u"He\U0000fffd\U0000fffdo".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding,
"replace")
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
# Test custom replacement character
test_value = np.array([72, 101, 2147483647, -1, 111], np.int32)
expected_value = u"Heooo".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding,
"replace", 111)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
# Verify "replace" is default
test_value = np.array([72, 101, 2147483647, -1, 111], np.int32)
expected_value = u"He\U0000fffd\U0000fffdo".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
# Replacement_char must be within range
test_value = np.array([72, 101, 2147483647, -1, 111], np.int32)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding,
"replace", 1114112)
- with self.cached_session():
- with self.assertRaises(errors.InvalidArgumentError):
- unicode_encode_op.eval()
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(unicode_encode_op)
# -- regular Tensor tests -- #
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def testVector(self, encoding):
test_value = np.array([72, 101, 108, 108, 111], np.int32)
expected_value = u"Hello".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
test_value = np.array([72, 101, 195, 195, 128516], np.int32)
expected_value = u"He\xc3\xc3\U0001f604".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
# Single character string
test_value = np.array([72], np.int32)
expected_value = u"H".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
test_value = np.array([128516], np.int32)
expected_value = u"\U0001f604".encode(encoding)
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(result, bytes)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def testMatrix(self, encoding):
test_value = np.array(
[[72, 128516, 108, 108, 111], [87, 128516, 114, 108, 100]], np.int32)
@@ -151,12 +145,10 @@
u"H\U0001f604llo".encode(encoding), u"W\U0001f604rld".encode(encoding)
]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(unicode_encode_op, ops.Tensor)
- self.assertAllEqual(result, expected_value)
+ self.assertAllEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def test3DimMatrix(self, encoding):
test_value = constant_op.constant(
[[[72, 101, 108, 108, 111], [87, 111, 114, 108, 100]],
@@ -166,12 +158,10 @@
[u"fixed".encode(encoding), u"words".encode(encoding)],
[u"Hyper".encode(encoding), u"cube.".encode(encoding)]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(unicode_encode_op, ops.Tensor)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def test4DimMatrix(self, encoding):
test_value = constant_op.constant(
[[[[72, 101, 108, 108, 111]], [[87, 111, 114, 108, 100]]],
@@ -184,14 +174,12 @@
[[u"Hyper".encode(encoding)],
[u"cube.".encode(encoding)]]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(unicode_encode_op, ops.Tensor)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
# -- Ragged Tensor tests -- #
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def testRaggedMatrix(self, encoding):
test_value = ragged_factory_ops.constant(
[[72, 195, 108, 108, 111], [87, 128516, 114, 108, 100, 46]], np.int32)
@@ -199,12 +187,10 @@
u"H\xc3llo".encode(encoding), u"W\U0001f604rld.".encode(encoding)
]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertIsInstance(unicode_encode_op, ops.Tensor)
- self.assertAllEqual(result, expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def test3DimMatrixWithRagged2ndDim(self, encoding):
test_value = ragged_factory_ops.constant(
[[[72, 101, 108, 108, 111], [87, 111, 114, 108, 100]],
@@ -218,12 +204,10 @@
u"cube.".encode(encoding)
]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertEqual(unicode_encode_op.ragged_rank, 1)
- self.assertAllEqual(result.tolist(), expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def test3DimMatrixWithRagged3rdDim(self, encoding):
test_value = ragged_factory_ops.constant(
[[[72, 101, 108, 108, 111], [87, 111, 114, 108, 100, 46]],
@@ -235,12 +219,10 @@
u"w\xc3rry, be".encode(encoding)
], [u"\U0001f604".encode(encoding), u"".encode(encoding)]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertEqual(unicode_encode_op.ragged_rank, 1)
- self.assertAllEqual(result.tolist(), expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def test3DimMatrixWithRagged2ndAnd3rdDim(self, encoding):
test_value = ragged_factory_ops.constant(
[[[72, 101, 108, 108, 111], [87, 111, 114, 108, 100, 46]], [],
@@ -248,12 +230,10 @@
expected_value = [[u"Hello".encode(encoding), u"World.".encode(encoding)],
[], [u"\U0001f604".encode(encoding)]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertEqual(unicode_encode_op.ragged_rank, 1)
- self.assertAllEqual(result.tolist(), expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def test4DimRaggedMatrix(self, encoding):
test_value = ragged_factory_ops.constant(
[[[[72, 101, 108, 108, 111], [87, 111, 114, 108, 100]]],
@@ -261,40 +241,30 @@
expected_value = [[[u"Hello".encode(encoding), u"World".encode(encoding)]],
[[u"".encode(encoding)], [u"Hype".encode(encoding)]]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertEqual(unicode_encode_op.ragged_rank, 2)
- self.assertAllEqual(result.tolist(), expected_value)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
@parameterized.parameters("UTF-8", "UTF-16-BE", "UTF-32-BE")
+ @test_util.run_v1_only("b/120545219")
def testRaggedMatrixWithMultiDimensionInnerValues(self, encoding):
- test_inner_values = constant_op.constant([[[72, 101, 108, 108, 111],
- [87, 111, 114, 108, 100]],
- [[102, 105, 120, 101, 100],
- [119, 111, 114, 100, 115]],
- [[72, 121, 112, 101, 114],
- [99, 117, 98, 101, 46]]])
+ test_flat_values = constant_op.constant([[[72, 101, 108, 108, 111],
+ [87, 111, 114, 108, 100]],
+ [[102, 105, 120, 101, 100],
+ [119, 111, 114, 100, 115]],
+ [[72, 121, 112, 101, 114],
+ [99, 117, 98, 101, 46]]])
test_row_splits = [
constant_op.constant([0, 2, 3], dtype=np.int64),
constant_op.constant([0, 1, 1, 3], dtype=np.int64)
]
- test_value = ragged_factory_ops.from_nested_row_splits(test_inner_values,
- test_row_splits)
+ test_value = ragged_tensor.RaggedTensor.from_nested_row_splits(
+ test_flat_values, test_row_splits)
expected_value = [[[[u"Hello".encode(encoding), u"World".encode(encoding)]],
[]],
[[[u"fixed".encode(encoding), u"words".encode(encoding)],
[u"Hyper".encode(encoding),
u"cube.".encode(encoding)]]]]
unicode_encode_op = ragged_string_ops.unicode_encode(test_value, encoding)
- with self.cached_session():
- result = unicode_encode_op.eval()
- self.assertEqual(unicode_encode_op.ragged_rank, 2)
- self.assertAllEqual(result.tolist(), expected_value)
- # These next two assertions don't necessarily need to be here as they test
- # internal representations and we already verified the value is correct.
- self.assertAllEqual(len(result.nested_row_splits), len(test_row_splits))
- self.assertEqual(unicode_encode_op.inner_values.shape.ndims,
- test_inner_values.shape.ndims - 1)
+ self.assertRaggedEqual(unicode_encode_op, expected_value)
if __name__ == "__main__":
diff --git a/tensorflow/python/kernel_tests/variable_scope_test.py b/tensorflow/python/kernel_tests/variable_scope_test.py
index 44d4bd5..451eb38 100644
--- a/tensorflow/python/kernel_tests/variable_scope_test.py
+++ b/tensorflow/python/kernel_tests/variable_scope_test.py
@@ -237,7 +237,8 @@
_ = d2(x)
self.assertEqual(len(d2.variables), 2)
v3, v4 = d2.variables
- self.assertAllEqual([v1, v2], [v3, v4])
+ self.assertEqual(v1, v3)
+ self.assertEqual(v2, v4)
f()
# TODO(mihaimaruseac): Not converted to use wrap_function because of
@@ -1684,7 +1685,7 @@
with variable_scope.variable_creator_scope(creator_b):
variable_scope.variable(1.0, name="one_name")
- self.assertAllEqual(variable_names, ["forced_name"])
+ self.assertEqual(variable_names[0], "forced_name")
called = [False]
diff --git a/tensorflow/python/kernel_tests/variables_test.py b/tensorflow/python/kernel_tests/variables_test.py
index 08d885e..07807e8 100644
--- a/tensorflow/python/kernel_tests/variables_test.py
+++ b/tensorflow/python/kernel_tests/variables_test.py
@@ -43,7 +43,7 @@
class VariablesTestCase(test.TestCase):
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitialization(self):
with self.cached_session():
var0 = variables.VariableV1(0.0)
@@ -66,12 +66,12 @@
with self.assertRaisesOpError("Attempting to use uninitialized value"):
self.evaluate(var1)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(0.0, self.evaluate(var0))
self.assertAllClose(1.1, self.evaluate(var1))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitializationOrder(self):
with self.cached_session():
rnd = variables.Variable(random_ops.random_uniform([3, 6]), name="rnd")
@@ -96,11 +96,11 @@
self.assertEqual([3, 6], depdep.get_shape())
self.assertEqual([3, 6], depdep.shape)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
- self.assertAllClose(rnd.eval(), self.evaluate(dep))
- self.assertAllClose(rnd.eval() + self.evaluate(dep) + 2.0,
- self.evaluate(depdep))
+ self.assertAllClose(self.evaluate(rnd), self.evaluate(dep))
+ self.assertAllClose(
+ self.evaluate(rnd) + self.evaluate(dep) + 2.0, self.evaluate(depdep))
def testIterable(self):
with self.assertRaisesRegexp(TypeError, "not iterable"):
@@ -117,7 +117,7 @@
plus_one = var.assign_add(1.0)
minus_one = var.assign_sub(2.0)
four = var.assign(4.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(0.0, self.evaluate(var))
self.assertAllClose(1.0, self.evaluate(plus_one))
@@ -136,7 +136,7 @@
plus_one = var.assign_add(1.0)
minus_one = var.assign_sub(2.0)
four = var.assign(4.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(0.0, self.evaluate(var))
self.evaluate(plus_one)
@@ -166,7 +166,7 @@
var = variables.Variable(zero)
count_up_to = var.count_up_to(3)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(0, self.evaluate(var))
self.assertEqual(0, self.evaluate(count_up_to))
@@ -194,7 +194,7 @@
def testCountUpToInt64(self):
self._countUpToTest(dtypes.int64)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testControlDepsNone(self):
with self.cached_session():
c = constant_op.constant(1.0)
@@ -208,7 +208,7 @@
self.assertEqual([], var_x.value().op.control_inputs)
self.assertEqual([], var_x._ref().op.control_inputs) # pylint: disable=protected-access
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testControlFlow(self):
with self.cached_session() as sess:
v0 = variables.Variable(0, name="v0")
@@ -245,7 +245,7 @@
self.evaluate(v0.initializer)
self.evaluate(add)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testControlFlowInitialization(self):
"""Expects an error if an initializer is in a control-flow scope."""
def cond(i, _):
@@ -264,10 +264,10 @@
with self.cached_session():
var_x = variables.Variable(2.0)
var_y = variables.Variable(3.0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(2.0, self.evaluate(var_x))
self.assertAllClose(3.0, self.evaluate(var_y))
- self.assertAllClose(5.0, math_ops.add(var_x, var_y).eval())
+ self.assertAllClose(5.0, self.evaluate(math_ops.add(var_x, var_y)))
@test_util.run_deprecated_v1
def testZeroSizeVarSameAsConst(self):
@@ -277,9 +277,9 @@
variable_mul = math_ops.matmul(zero_size_const, zero_size_var)
const_mul = math_ops.matmul(
zero_size_const, zero_size_const, transpose_b=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variable_output = self.evaluate(variable_mul)
- self.assertAllClose(const_mul.eval(), variable_output)
+ self.assertAllClose(self.evaluate(const_mul), variable_output)
self.assertAllClose([[0., 0.], [0., 0.]], variable_output)
@test_util.run_deprecated_v1
@@ -372,7 +372,7 @@
matmul = var_m.__matmul__([[10.0], [20.0]])
rmatmul = var_m.__rmatmul__([[10.0], [20.0]])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose([2.0], self.evaluate(add))
self.assertAllClose([3.0], self.evaluate(radd))
self.assertAllClose([1.0], self.evaluate(sub))
@@ -409,10 +409,10 @@
def testSession(self):
with self.cached_session() as sess:
var = variables.Variable([1, 12])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose([1, 12], self.evaluate(var))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testColocation(self):
with ops.device("/job:ps"):
var = variables.VariableV1(0, name="v")
@@ -421,7 +421,7 @@
self.assertDeviceEqual("/job:ps", assign_op.device)
self.assertEqual([b"loc:@v"], assign_op.op.colocation_groups())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitializerFunction(self):
value = [[-42], [133.7]]
shape = [2, 1]
@@ -431,7 +431,7 @@
v1 = variables.Variable(initializer, dtype=dtypes.float32)
self.assertEqual(shape, v1.get_shape())
self.assertEqual(shape, v1.shape)
- self.assertAllClose(value, v1.initial_value.eval())
+ self.assertAllClose(value, self.evaluate(v1.initial_value))
with self.assertRaises(errors_impl.FailedPreconditionError):
self.evaluate(v1)
@@ -439,11 +439,11 @@
math_ops.negative(v1.initialized_value()), dtype=dtypes.float32)
self.assertEqual(v1.get_shape(), v2.get_shape())
self.assertEqual(v1.shape, v2.shape)
- self.assertAllClose(np.negative(value), v2.initial_value.eval())
+ self.assertAllClose(np.negative(value), self.evaluate(v2.initial_value))
with self.assertRaises(errors_impl.FailedPreconditionError):
self.evaluate(v2)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(np.negative(value), self.evaluate(v2))
def testConstraintArg(self):
@@ -459,16 +459,16 @@
lambda: constant_op.constant(1.),
constraint=constraint)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testNoRefDataRace(self):
with self.cached_session():
a = variables.Variable([1, 2, 3], dtype=dtypes.float32)
b = variables.Variable(a.initialized_value() + 2)
c = variables.Variable(b.initialized_value() + 2)
- variables.global_variables_initializer().run()
- self.assertAllEqual(a.eval(), [1, 2, 3])
- self.assertAllEqual(b.eval(), [3, 4, 5])
- self.assertAllEqual(c.eval(), [5, 6, 7])
+ self.evaluate(variables.global_variables_initializer())
+ self.assertAllEqual(self.evaluate(a), [1, 2, 3])
+ self.assertAllEqual(self.evaluate(b), [3, 4, 5])
+ self.assertAllEqual(self.evaluate(c), [5, 6, 7])
@test_util.run_deprecated_v1
def testInitializerFunctionDevicePlacement(self):
@@ -489,7 +489,7 @@
for i in v2.initializer.inputs:
self.assertEqual(expected_group_v2, i.op.colocation_groups())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testVariableDefInitializedInstances(self):
with ops.Graph().as_default(), self.cached_session() as sess:
v_def = variables.Variable(
@@ -503,7 +503,7 @@
# initialized_value should not rerun the initializer_op if the variable
# has already been initialized elsewhere.
self.evaluate(v.assign(1.0))
- self.assertEqual(1.0, v.initialized_value().eval())
+ self.assertEqual(1.0, self.evaluate(v.initialized_value()))
v_def.ClearField("initial_value_name")
with ops.Graph().as_default(), self.cached_session() as sess:
@@ -537,12 +537,12 @@
def testLoad(self):
with self.cached_session():
var = variables.Variable(np.zeros((5, 5), np.float32))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
var.load(np.ones((5, 5), np.float32))
self.assertAllClose(np.ones((5, 5), np.float32), self.evaluate(var))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRepr(self):
var = variables.VariableV1(np.zeros((5, 5), np.float32), name="noop")
self.assertEqual(
@@ -573,10 +573,10 @@
_ = v, w
uninited = variables.report_uninitialized_variables()
self.assertAllEqual(np.array([b"v", b"w"]), self.evaluate(uninited))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual(0, self.evaluate(uninited).size)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testVariableList(self):
with ops.Graph().as_default(), self.cached_session() as sess:
v = variables.VariableV1([1, 2], name="v")
@@ -601,20 +601,20 @@
b = variables.Variable(array_ops.ones([2, 2]))
objective = math_ops.reduce_sum(b + math_ops.matmul(
a, a, transpose_a=True))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
do_opt = gradient_descent.GradientDescentOptimizer(0.1).minimize(
objective)
self.evaluate([do_opt])
self.assertAllClose([[0.9, 0.9], [0.9, 0.9]], self.evaluate(b))
+@test_util.run_v1_only("b/120545219")
class ObsoleteIsInitializedTest(test.TestCase):
def testNoVars(self):
with ops.Graph().as_default():
self.assertEqual(None, variables.assert_variables_initialized())
- @test_util.run_deprecated_v1
def testVariables(self):
with ops.Graph().as_default(), self.cached_session() as sess:
v = variables.VariableV1([1, 2])
@@ -623,10 +623,9 @@
inited = variables.assert_variables_initialized()
with self.assertRaisesOpError("Attempting to use uninitialized value"):
self.evaluate(inited)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.evaluate(inited)
- @test_util.run_deprecated_v1
def testVariableList(self):
with ops.Graph().as_default(), self.cached_session() as sess:
v = variables.VariableV1([1, 2])
@@ -766,36 +765,36 @@
assign_list = pv_1.assign([c_0, c_1])
assign_part_value = pv_1.assign_add(assign_ones)
assign_part_var = pv_1.assign_sub(pv_0)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
- self.assertEqual([1.0], plus_delta[0].eval())
+ self.assertEqual([1.0], self.evaluate(plus_delta[0]))
self.assertEqual([1.0], self.evaluate(v0))
- self.assertEqual([3.0], plus_delta[1].eval())
+ self.assertEqual([3.0], self.evaluate(plus_delta[1]))
self.assertEqual([3.0], self.evaluate(v1))
- self.assertEqual([-2.0], minus_delta[0].eval())
+ self.assertEqual([-2.0], self.evaluate(minus_delta[0]))
self.assertEqual([-2.0], self.evaluate(v0))
- self.assertEqual([-1.0], minus_delta[1].eval())
+ self.assertEqual([-1.0], self.evaluate(minus_delta[1]))
self.assertEqual([-1.0], self.evaluate(v1))
- self.assertEqual([1.0], assign_ones[0].eval())
+ self.assertEqual([1.0], self.evaluate(assign_ones[0]))
self.assertEqual([1.0], self.evaluate(v0))
- self.assertEqual([1.0], assign_ones[1].eval())
+ self.assertEqual([1.0], self.evaluate(assign_ones[1]))
self.assertEqual([1.0], self.evaluate(v1))
- self.assertEqual([2.0], assign_list[0].eval())
+ self.assertEqual([2.0], self.evaluate(assign_list[0]))
self.assertEqual([2.0], self.evaluate(v2))
- self.assertEqual([3.0], assign_list[1].eval())
+ self.assertEqual([3.0], self.evaluate(assign_list[1]))
self.assertEqual([3.0], self.evaluate(v3))
- self.assertEqual([3.0], assign_part_value[0].eval())
+ self.assertEqual([3.0], self.evaluate(assign_part_value[0]))
self.assertEqual([3.0], self.evaluate(v2))
- self.assertEqual([4.0], assign_part_value[1].eval())
+ self.assertEqual([4.0], self.evaluate(assign_part_value[1]))
self.assertEqual([4.0], self.evaluate(v3))
- self.assertEqual([2.0], assign_part_var[0].eval())
+ self.assertEqual([2.0], self.evaluate(assign_part_var[0]))
self.assertEqual([2.0], self.evaluate(v2))
- self.assertEqual([3.0], assign_part_var[1].eval())
+ self.assertEqual([3.0], self.evaluate(assign_part_var[1]))
self.assertEqual([3.0], self.evaluate(v3))
diff --git a/tensorflow/python/kernel_tests/while_v2_test.py b/tensorflow/python/kernel_tests/while_v2_test.py
index 09cbeb1..cae459a 100644
--- a/tensorflow/python/kernel_tests/while_v2_test.py
+++ b/tensorflow/python/kernel_tests/while_v2_test.py
@@ -53,6 +53,7 @@
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
+ @test_util.run_v1_only("b/120545219")
def testReturnSameStructureTrue(self):
x = constant_op.constant(2.)
ret = while_loop_v2(
@@ -145,7 +146,7 @@
self.assertSequenceEqual(self.evaluate(grad), [32.])
self.assertSequenceEqual(self.evaluate(grad_grad), [48.])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPruning(self):
x = constant_op.constant(1)
@@ -252,7 +253,7 @@
if op.type == "While":
while_op = op
- body_graph = while_v2._get_body_graph(while_op)
+ body_graph = while_v2._get_graph(while_op, "body")
# body_graph.inputs: [counter_arg, x_arg, tl_arg, *accumulators]
x_input_t = body_graph.inputs[1]
accumulator_count = len(
@@ -283,22 +284,26 @@
self.assertListEqual(actual_tensor_shape.as_list(), shape)
def GetAccumulatorForInputAtIndex(while_op, idx):
- body_graph = while_v2._get_body_graph(while_op)
+ body_graph = while_v2._get_graph(while_op, "body")
y_input_t = body_graph.inputs[idx]
push_back_node = [c for c in y_input_t.consumers()
if c.type == "TensorListPushBack"][0]
output_idx = body_graph.outputs.index(push_back_node.outputs[0])
return while_op.outputs[output_idx]
- x = constant_op.constant(2.)
+ x = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
y = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
# Forward pass.
- ret = while_loop_v2(
- lambda v, u: v < 8.,
- lambda v, u: (v * v, u), [x, y],
- return_same_structure=False)
+ ret = while_loop_v2(lambda v, u: v < 8.,
+ lambda v, u: (math_ops.pow(v, u), u),
+ [x, y],
+ return_same_structure=True)
while_op = ret[0].op.inputs[0].op
+ # Gradient pass.
+ grad = gradients_impl.gradients(ret[0], x)
+ grad_while_op = grad[0].op.inputs[0].op
+
# Get the TensorList output of While op containing the accumulated values
# of y.
# while_op.inputs: [counter_arg, x_arg, y_arg, *accumulators]
@@ -307,14 +312,14 @@
element_dtype=dtypes.float32)
MatchShape(val.shape)
- # Gradient pass.
- grad = gradients_impl.gradients(ret[1], y)
- grad_while_op = grad[0].op.inputs[0].op
+ # Take second derivative to generate intermediate grad_while_op outputs
+ gradients_impl.gradients(grad, x)
+
# Get the TensorList output of gradient While op containing the accumulated
- # values of grad_y.
+ # values of grad_x (note that grad_x is needed by the second derivative).
# grad_while_op.inputs:
# [counter_arg, total_iters_arg, grad_x_arg, grad_y_arg, *other_args]
- grad_output = GetAccumulatorForInputAtIndex(grad_while_op, 3)
+ grad_output = GetAccumulatorForInputAtIndex(grad_while_op, 2)
_, val = list_ops.tensor_list_pop_back(grad_output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
@@ -357,14 +362,13 @@
@test_util.enable_control_flow_v2
@test_util.run_deprecated_v1
def testWhileAndTensorArray(self):
- with self.cached_session() as sess:
- param = constant_op.constant(2.0)
- y0 = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="elems")
- # map_fn uses TensorArray internally.
- r = functional_ops.map_fn(lambda x: math_ops.multiply(x, param), y0)
- self.assertAllClose([2.0, 4.0, 6.0, 8.0, 10.0, 12.0], self.evaluate(r))
- r = gradients_impl.gradients(r, param)[0]
- self.assertAllClose(21.0, self.evaluate(r))
+ param = constant_op.constant(2.0)
+ y0 = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="elems")
+ # map_fn uses TensorArray internally.
+ r = functional_ops.map_fn(lambda x: math_ops.multiply(x, param), y0)
+ grad = gradients_impl.gradients(r, param)[0]
+ self.assertAllClose([2.0, 4.0, 6.0, 8.0, 10.0, 12.0], self.evaluate(r))
+ self.assertAllClose(21.0, self.evaluate(grad))
@test_util.run_deprecated_v1
def testNestedWhile(self):
@@ -438,6 +442,26 @@
# grad = gradients_impl.gradients(output, [n])
# self.assertEqual(self.evaluate(grad), 3.5)
+ @test_util.run_deprecated_v1
+ def testForwardPassRewrite(self):
+ x = constant_op.constant(1.0, name="x")
+ output = while_v2.while_loop(lambda x: x < 10.0,
+ lambda x: x * 2.0,
+ [x])[0]
+ while_op = output.op.inputs[0].op
+ self.assertEqual(while_op.type, "While")
+ # outputs = [loop_counter, x]
+ self.assertLen(while_op.outputs, 2)
+
+ gradients_impl.gradients(output, x)
+ # while_op should have been rewritten to output 2.0 intermediate.
+ # outputs = [loop_counter, x, 2.0_accumulator, x_accumulator]
+ self.assertLen(while_op.outputs, 4)
+
+ gradients_impl.gradients(output, x)
+ # Computing the gradient again shouldn't rewrite while_op again.
+ self.assertLen(while_op.outputs, 4)
+
def ScalarShape():
return ops.convert_to_tensor([], dtype=dtypes.int32)
diff --git a/tensorflow/python/layers/core_test.py b/tensorflow/python/layers/core_test.py
index cf6f0fb..b40a268 100644
--- a/tensorflow/python/layers/core_test.py
+++ b/tensorflow/python/layers/core_test.py
@@ -533,10 +533,13 @@
self.assertEqual(y.get_shape().as_list(), [None, 6])
@test_util.run_deprecated_v1
- def testFlattenValueError(self):
+ def testFlatten0D(self):
x = array_ops.placeholder(shape=(None,), dtype='float32')
- with self.assertRaises(ValueError):
- core_layers.Flatten()(x)
+ y = core_layers.Flatten()(x)
+ with self.cached_session() as sess:
+ np_output = sess.run(y, feed_dict={x: np.zeros((5,))})
+ self.assertEqual(list(np_output.shape), [5, 1])
+ self.assertEqual(y.shape.as_list(), [None, 1])
@test_util.run_deprecated_v1
def testFlattenUnknownAxes(self):
diff --git a/tensorflow/python/layers/normalization_test.py b/tensorflow/python/layers/normalization_test.py
index 07d8e40..6535f74 100644
--- a/tensorflow/python/layers/normalization_test.py
+++ b/tensorflow/python/layers/normalization_test.py
@@ -38,6 +38,7 @@
from tensorflow.python.training import saver as saver_lib
+@test_util.run_v1_only('b/120545219')
class BNTest(test.TestCase):
def _simple_model(self, image, fused, freeze_mode):
@@ -144,7 +145,6 @@
return train_vars, loss_val
- @test_util.run_deprecated_v1
def testHalfPrecision(self):
ref_vars, ref_loss = self._trainEvalSequence(
dtype=dtypes.float32,
@@ -230,43 +230,33 @@
ckpt_b_use_gpu, use_gpu_test_a, use_gpu_test_b,
freeze_mode)
- @test_util.run_deprecated_v1
def testCheckpointFusedCPUAndFusedGPU(self):
self._testCheckpointCrossDevice(True, False, True, True)
- @test_util.run_deprecated_v1
def testCheckpointFusedCPUAndFusedCPU(self):
self._testCheckpointCrossDevice(True, False, True, False)
- @test_util.run_deprecated_v1
def testCheckpointFusedGPUAndFusedGPU(self):
self._testCheckpointCrossDevice(True, True, True, True)
- @test_util.run_deprecated_v1
def testCheckpointNonFusedCPUAndNonFusedGPU(self):
self._testCheckpointCrossDevice(False, False, False, True)
- @test_util.run_deprecated_v1
def testCheckpointNonFusedCPUAndNonFusedCPU(self):
self._testCheckpointCrossDevice(False, False, False, False)
- @test_util.run_deprecated_v1
def testCheckpointNonFusedGPUAndNonFusedGPU(self):
self._testCheckpointCrossDevice(False, True, False, True)
- @test_util.run_deprecated_v1
def testCheckpointNonFusedGPUAndFusedGPU(self):
self._testCheckpointCrossDevice(False, True, True, True)
- @test_util.run_deprecated_v1
def testCheckpointNonFusedGPUAndFusedCPU(self):
self._testCheckpointCrossDevice(False, True, True, False)
- @test_util.run_deprecated_v1
def testCheckpointNonFusedCPUAndFusedCPU(self):
self._testCheckpointCrossDevice(False, False, True, False)
- @test_util.run_deprecated_v1
def testCreateBN(self):
# Call layer.
bn = normalization_layers.BatchNormalization(axis=1)
@@ -293,7 +283,6 @@
ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES),
bn.trainable_variables)
- @test_util.run_deprecated_v1
def testCreateFusedBNFloat16(self):
# Call layer.
bn = normalization_layers.BatchNormalization(axis=1, fused=True)
@@ -323,7 +312,6 @@
ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES),
bn.trainable_variables)
- @test_util.run_deprecated_v1
def test3DInputAxis1(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -367,7 +355,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def test3DInputAxis2(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -451,7 +438,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def test4DInputAxis2(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -493,7 +479,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def test4DInputAxis3(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -535,7 +520,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def test4DInputAxis3Fused(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -619,7 +603,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def testNegativeAxis(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -662,7 +645,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def testBooleanLearningPhase(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -703,7 +685,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def testFunctionalNoReuse(self):
inputs = variables.Variable(
np.random.random((5, 4, 3, 6)), dtype=dtypes.float32)
@@ -756,7 +737,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def testFunctionalReuse(self):
inputs1 = variables.Variable(
np.random.random((5, 4, 3, 6)), dtype=dtypes.float32)
@@ -821,7 +801,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=2)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def testFunctionalReuseFromScope(self):
inputs = variables.Variable(
np.random.random((5, 4, 3, 6)), dtype=dtypes.float32)
@@ -836,7 +815,6 @@
inputs, axis=-1, momentum=0.9, epsilon=epsilon, training=training)
self.assertEqual(len(variables.global_variables()), 5)
- @test_util.run_deprecated_v1
def testNoCenter(self):
bn = normalization_layers.BatchNormalization(axis=1, center=False)
inputs = random_ops.random_uniform((5, 4, 3), seed=1)
@@ -852,7 +830,6 @@
self.assertEqual(len(bn.trainable_variables), 1)
self.assertEqual(len(bn.non_trainable_variables), 2)
- @test_util.run_deprecated_v1
def testNoScale(self):
bn = normalization_layers.BatchNormalization(axis=1, scale=False)
inputs = random_ops.random_uniform((5, 4, 3), seed=1)
@@ -868,7 +845,6 @@
self.assertEqual(len(bn.trainable_variables), 1)
self.assertEqual(len(bn.non_trainable_variables), 2)
- @test_util.run_deprecated_v1
def testRegularizers(self):
reg = lambda x: 0.1 * math_ops.reduce_sum(x)
bn = normalization_layers.BatchNormalization(axis=1, beta_regularizer=reg)
@@ -894,7 +870,6 @@
self.assertEqual(bn.gamma_constraint, g_constraint)
self.assertEqual(bn.beta_constraint, b_constraint)
- @test_util.run_deprecated_v1
def testRenorm(self):
shape = (4, 3)
xt = array_ops.placeholder(dtypes.float32, shape)
@@ -953,7 +928,6 @@
self.assertAllClose(y_train, yt_val_train, atol=1e-5)
self.assertAllClose(y_test, yt_val_test, atol=1e-5)
- @test_util.run_deprecated_v1
def testAdjustment(self):
shape = (4, 3)
xt = array_ops.placeholder(dtypes.float32, shape)
@@ -998,7 +972,6 @@
self.assertAllClose(y_train, yt_val_train, atol=1e-5)
self.assertAllClose(y_test, yt_val_test, atol=1e-5)
- @test_util.run_deprecated_v1
def testRenormWithAdjustment(self):
shape = (4, 3)
xt = array_ops.placeholder(dtypes.float32, shape)
@@ -1069,7 +1042,6 @@
normalization_layers.batch_normalization(
inp, virtual_batch_size=-1)
- @test_util.run_deprecated_v1
def testGhostBNVirtualBatchFull(self):
shape = [6, 5, 4, 3]
inp = random_ops.random_uniform(shape, seed=1)
@@ -1095,7 +1067,6 @@
inp, virtual_batch_size=3)
self.assertListEqual(out.shape.as_list(), shape)
- @test_util.run_deprecated_v1
def testGhostBNUnknownBatchSize(self):
np_shape = [10, 5, 4]
tf_shape = [None, 5, 4]
@@ -1111,7 +1082,6 @@
self.assertListEqual(list(y.shape), np_shape)
- @test_util.run_deprecated_v1
def testGhostBN2Dims(self):
shape = [6, 2]
virtual_batch_size = 3
@@ -1165,7 +1135,6 @@
self.assertAllClose(y_train, y_val_train, atol=1e-5)
self.assertAllClose(y_test, y_val_test, atol=1e-5)
- @test_util.run_deprecated_v1
def testGhostBN4DimsAxis3(self):
shape = [6, 10, 10, 3]
virtual_batch_size = 2
@@ -1219,7 +1188,6 @@
self.assertAllClose(y_train, y_val_train, atol=1e-2)
self.assertAllClose(y_test, y_val_test, atol=1e-2)
- @test_util.run_deprecated_v1
def testGhostBN4DimsAxis1(self):
shape = [6, 3, 10, 10]
virtual_batch_size = 2
@@ -1290,7 +1258,6 @@
normalization_layers.batch_normalization(
inp, axis=[1, 2, 1]) # duplicate
- @test_util.run_deprecated_v1
def test3DInputMultiAxis12(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -1332,7 +1299,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def test5DInputMultiAxis123(self):
epsilon = 1e-3
bn = normalization_layers.BatchNormalization(
@@ -1374,7 +1340,6 @@
self.assertAlmostEqual(np.mean(normed_np_output), 0., places=1)
self.assertAlmostEqual(np.std(normed_np_output), 1., places=1)
- @test_util.run_deprecated_v1
def testGhostBN5DimsMultiAxis14(self):
shape = [6, 3, 10, 10, 4]
virtual_batch_size = 3
diff --git a/tensorflow/python/lib/io/file_io.py b/tensorflow/python/lib/io/file_io.py
index 4caa575..ee55d89 100644
--- a/tensorflow/python/lib/io/file_io.py
+++ b/tensorflow/python/lib/io/file_io.py
@@ -671,7 +671,7 @@
@tf_export("io.gfile.walk")
-def walk_v2(top, topdown, onerror=None):
+def walk_v2(top, topdown=True, onerror=None):
"""Recursive directory tree generator for directories.
Args:
diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
index 7b6242b..e10d903 100644
--- a/tensorflow/python/ops/array_ops.py
+++ b/tensorflow/python/ops/array_ops.py
@@ -56,6 +56,7 @@
@tf_export("identity")
+@dispatch.add_dispatch_support
def identity(input, name=None): # pylint: disable=redefined-builtin
r"""Return a tensor with the same shape and contents as input.
@@ -77,11 +78,16 @@
return input._copy() # pylint: disable=protected-access
return input
else:
- return gen_array_ops.identity(input, name=name)
+ ret = gen_array_ops.identity(input, name=name)
+ # Propagate handle data for happier shape inference for resource variables.
+ if hasattr(input, "_handle_data"):
+ ret._handle_data = input._handle_data # pylint: disable=protected-access
+ return ret
# pylint: disable=redefined-builtin,protected-access
@tf_export(v1=["expand_dims"])
+@dispatch.add_dispatch_support
@deprecation.deprecated_args(None, "Use the `axis` argument instead", "dim")
def expand_dims(input, axis=None, name=None, dim=None):
"""Inserts a dimension of 1 into a tensor's shape.
@@ -139,6 +145,7 @@
@tf_export("expand_dims", v1=[])
+@dispatch.add_dispatch_support
def expand_dims_v2(input, axis, name=None):
"""Inserts a dimension of 1 into a tensor's shape.
@@ -941,6 +948,7 @@
@tf_export("stack")
+@dispatch.add_dispatch_support
def stack(values, axis=0, name="stack"):
"""Stacks a list of rank-`R` tensors into one rank-`(R+1)` tensor.
@@ -1151,6 +1159,7 @@
@tf_export("concat")
+@dispatch.add_dispatch_support
def concat(values, axis, name="concat"):
"""Concatenates tensors along one dimension.
@@ -1328,6 +1337,7 @@
@tf_export("boolean_mask", v1=[])
+@dispatch.add_dispatch_support
def boolean_mask_v2(tensor, mask, axis=None, name="boolean_mask"):
"""Apply boolean mask to tensor.
@@ -1810,6 +1820,7 @@
@tf_export(v1=["zeros_like"])
+@dispatch.add_dispatch_support
def zeros_like(tensor, dtype=None, name=None, optimize=True):
"""Creates a tensor with all elements set to zero.
@@ -1840,6 +1851,7 @@
@tf_export("zeros_like", v1=[])
+@dispatch.add_dispatch_support
def zeros_like_v2(
input, # pylint: disable=redefined-builtin
dtype=None,
@@ -1899,6 +1911,7 @@
@tf_export(v1=["ones_like"])
+@dispatch.add_dispatch_support
def ones_like(tensor, dtype=None, name=None, optimize=True):
"""Creates a tensor with all elements set to 1.
@@ -1929,6 +1942,7 @@
@tf_export("ones_like", v1=[])
+@dispatch.add_dispatch_support
def ones_like_v2(
input, # pylint: disable=redefined-builtin
dtype=None,
@@ -2643,7 +2657,7 @@
return result_paddings, result_crops
-@tf_export("nn.space_to_batch", v1=["nn.space_to_batch", "space_to_batch"])
+@tf_export(v1=["nn.space_to_batch", "space_to_batch"])
@deprecation.deprecated_endpoints("space_to_batch")
def space_to_batch(input, paddings, block_size, name=None): # pylint: disable=redefined-builtin
result = space_to_batch_nd(
@@ -2658,7 +2672,15 @@
space_to_batch.__doc__ = gen_array_ops.space_to_batch.__doc__
-@tf_export("nn.space_to_depth", v1=["nn.space_to_depth", "space_to_depth"])
+@tf_export("space_to_batch", "nn.space_to_batch", v1=[])
+def space_to_batch_v2(input, block_shape, paddings, name=None): # pylint: disable=redefined-builtin
+ return space_to_batch_nd(input, block_shape, paddings, name)
+
+
+space_to_batch_v2.__doc__ = gen_array_ops.space_to_batch_nd.__doc__
+
+
+@tf_export(v1=["nn.space_to_depth", "space_to_depth"])
@deprecation.deprecated_endpoints("space_to_depth")
def space_to_depth(input, block_size, name=None, data_format="NHWC"): # pylint: disable=redefined-builtin
return gen_array_ops.space_to_depth(input, block_size, data_format, name=name)
@@ -2667,7 +2689,15 @@
space_to_depth.__doc__ = gen_array_ops.space_to_depth.__doc__
-@tf_export("nn.depth_to_space", v1=["nn.depth_to_space", "depth_to_space"])
+@tf_export("nn.space_to_depth", v1=[])
+def space_to_depth_v2(input, block_size, data_format="NHWC", name=None): # pylint: disable=redefined-builtin
+ return gen_array_ops.space_to_depth(input, block_size, data_format, name=name)
+
+
+space_to_depth_v2.__doc__ = gen_array_ops.space_to_depth.__doc__
+
+
+@tf_export(v1=["nn.depth_to_space", "depth_to_space"])
@deprecation.deprecated_endpoints("depth_to_space")
def depth_to_space(input, block_size, name=None, data_format="NHWC"): # pylint: disable=redefined-builtin
return gen_array_ops.depth_to_space(input, block_size, data_format, name=name)
@@ -2676,6 +2706,14 @@
depth_to_space.__doc__ = gen_array_ops.depth_to_space.__doc__
+@tf_export("nn.depth_to_space", v1=[])
+def depth_to_space_v2(input, block_size, data_format="NHWC", name=None): # pylint: disable=redefined-builtin
+ return gen_array_ops.depth_to_space(input, block_size, data_format, name=name)
+
+
+depth_to_space_v2.__doc__ = gen_array_ops.depth_to_space.__doc__
+
+
@tf_export(v1=["batch_to_space"])
def batch_to_space(input, crops, block_size, name=None): # pylint: disable=redefined-builtin
result = batch_to_space_nd(
@@ -3115,6 +3153,7 @@
@tf_export("where")
+@dispatch.add_dispatch_support
def where(condition, x=None, y=None, name=None):
"""Return the elements, either from `x` or `y`, depending on the `condition`.
@@ -3218,6 +3257,7 @@
@tf_export(v1=["gather"])
+@dispatch.add_dispatch_support
def gather(params, indices, validate_indices=None, name=None, axis=0):
del validate_indices
if axis != 0:
@@ -3234,6 +3274,7 @@
@tf_export("gather", v1=[])
+@dispatch.add_dispatch_support
def gather_v2(params, indices, validate_indices=None, axis=0, name=None):
return gather(params, indices, validate_indices=validate_indices, name=name,
axis=axis)
diff --git a/tensorflow/python/ops/candidate_sampling_ops.py b/tensorflow/python/ops/candidate_sampling_ops.py
index c64000b..56f76a4 100644
--- a/tensorflow/python/ops/candidate_sampling_ops.py
+++ b/tensorflow/python/ops/candidate_sampling_ops.py
@@ -151,7 +151,10 @@
seed2=seed2, name=name)
-@tf_export('nn.learned_unigram_candidate_sampler')
+@tf_export(
+ 'random.learned_unigram_candidate_sampler',
+ 'nn.learned_unigram_candidate_sampler')
+@deprecation.deprecated_endpoints(['nn.learned_unigram_candidate_sampler'])
def learned_unigram_candidate_sampler(true_classes, num_true, num_sampled,
unique, range_max, seed=None, name=None):
"""Samples a set of classes from a distribution learned during training.
@@ -209,8 +212,7 @@
@tf_export('random.fixed_unigram_candidate_sampler',
- 'nn.fixed_unigram_candidate_sampler',
- v1=['nn.fixed_unigram_candidate_sampler'])
+ 'nn.fixed_unigram_candidate_sampler')
def fixed_unigram_candidate_sampler(true_classes,
num_true,
num_sampled,
@@ -302,8 +304,7 @@
unigrams=unigrams, seed=seed1, seed2=seed2, name=name)
-@tf_export('random.all_candidate_sampler', 'nn.all_candidate_sampler',
- v1=['nn.all_candidate_sampler'])
+@tf_export('random.all_candidate_sampler', 'nn.all_candidate_sampler')
def all_candidate_sampler(true_classes, num_true, num_sampled, unique,
seed=None, name=None):
"""Generate the set of all classes.
diff --git a/tensorflow/python/ops/clip_ops.py b/tensorflow/python/ops/clip_ops.py
index 82803ac..a237cff 100644
--- a/tensorflow/python/ops/clip_ops.py
+++ b/tensorflow/python/ops/clip_ops.py
@@ -31,10 +31,12 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import numerics
from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
from tensorflow.python.util.tf_export import tf_export
@tf_export("clip_by_value")
+@dispatch.add_dispatch_support
def clip_by_value(t, clip_value_min, clip_value_max,
name=None):
"""Clips tensor values to a specified min and max.
diff --git a/tensorflow/python/ops/cond_v2.py b/tensorflow/python/ops/cond_v2.py
index f7a95bd..7d09e32 100644
--- a/tensorflow/python/ops/cond_v2.py
+++ b/tensorflow/python/ops/cond_v2.py
@@ -61,7 +61,7 @@
# Automatic control dependencies are added in defuns, but not in v1
# graphs. Propagate that behavior here.
- add_control_dependencies = util.in_defun()
+ add_control_dependencies = ops.get_default_graph()._add_control_dependencies
pred = ops.convert_to_tensor(pred)
true_graph = func_graph_module.func_graph_from_py_func(
@@ -108,6 +108,46 @@
assert ([t.dtype for t in true_grad_graph.outputs] ==
[t.dtype for t in false_grad_graph.outputs])
+ if (true_grad_graph.if_op_needs_rewrite or
+ false_grad_graph.if_op_needs_rewrite):
+ # Modify 'op' to output the intermediates needed by the grad functions. Note
+ # that all needed intermediates are wrapped in optionals. Each optional
+ # intermediate output will have a value iff its corresponding branch is
+ # taken.
+ # NOTE(skyewm): if there are any active sessions, this modification to `op`
+ # may make them unrunnable!
+
+ if control_flow_util.InXlaContext(ops.get_default_graph()):
+ # XLA does not yet support optionals, so output intermediates directly and
+ # make them match via FakeParams, which can be converted to zeros in XLA.
+ # TODO(skyewm,jpienaar): can XLA support optionals?
+ true_intermediates = true_grad_graph.xla_intermediates
+ false_intermediates = false_grad_graph.xla_intermediates
+ extra_true_outputs, extra_false_outputs = _make_intermediates_match_xla(
+ true_graph, false_graph, true_intermediates, false_intermediates)
+ else:
+ true_intermediates = true_grad_graph.wrapped_intermediates
+ false_intermediates = false_grad_graph.wrapped_intermediates
+ # Make outputs match by adding none optionals.
+ extra_true_outputs, extra_false_outputs = _make_intermediates_match(
+ true_graph, false_graph, true_intermediates, false_intermediates)
+
+ true_graph.outputs.extend(extra_true_outputs)
+ false_graph.outputs.extend(extra_false_outputs)
+ # TODO(skyewm): indicate it's an internal bug if this fails.
+ _check_same_outputs(true_graph, false_graph)
+
+ true_graph.name += "_rewritten"
+ false_graph.name += "_rewritten"
+
+ op._set_func_attr("then_branch", util.create_new_tf_function(true_graph))
+ op._set_func_attr("else_branch", util.create_new_tf_function(false_graph))
+ op._set_type_list_attr("Tout", true_graph.output_types)
+ op._set_shape_list_attr("output_shapes", true_graph.output_shapes)
+ op._add_outputs(
+ [t.dtype for t in extra_true_outputs],
+ [t.shape for t in extra_true_outputs])
+
# Resolve references to forward graph tensors in grad graphs and ensure
# they are in-scope, i.e., belong to one of outer graphs of the grad graph.
true_grad_inputs = _resolve_grad_inputs(true_graph, true_grad_graph)
@@ -150,40 +190,6 @@
cond_inputs = _make_inputs_match(true_graph, false_graph,
true_inputs, false_inputs)
- # Add all intermediate tensors as function outputs so they're available for
- # the gradient computation. Since the outputs of the two functions must match,
- # we wrap all the intermediates in optionals. Each intermediate output will
- # have a value iff its corresponding branch is taken.
-
- true_intermediates = _get_intermediates(true_graph)
- false_intermediates = _get_intermediates(false_graph)
-
- # Save the original number of outputs to return to the caller.
- num_cond_outputs = len(true_graph.outputs)
-
- if control_flow_util.InXlaContext(ops.get_default_graph()):
- # XLA does not yet support optionals, so output intermediates directly and
- # make them match via FakeParams, which can be converted to zeros in XLA.
- # TODO(skyewm,jpienaar): can XLA support optionals?
- extra_true_outputs, extra_false_outputs = _make_intermediates_match_xla(
- true_graph, false_graph, true_intermediates, false_intermediates)
- else:
- # Wrap intermediates in optionals.
- wrapped_true_intermediates = _wrap_intermediates(true_graph,
- true_intermediates)
- wrapped_false_intermediates = _wrap_intermediates(false_graph,
- false_intermediates)
-
- # Make outputs match by adding none optionals.
- extra_true_outputs, extra_false_outputs = _make_intermediates_match(
- true_graph, false_graph,
- wrapped_true_intermediates, wrapped_false_intermediates)
-
- true_graph.outputs.extend(extra_true_outputs)
- false_graph.outputs.extend(extra_false_outputs)
- # TODO(skyewm): somehow indicate it's a bug if this fails.
- _check_same_outputs(true_graph, false_graph)
-
# Create the If op.
tensors = gen_functional_ops._if( # pylint: disable=protected-access
pred,
@@ -210,8 +216,7 @@
# Prevent fetching since the variant outputs can't be fetched directly.
if_op.graph.prevent_fetching(if_op)
-
- return tensors[:num_cond_outputs]
+ return tensors
def _get_func_graphs(if_op):
@@ -540,13 +545,37 @@
class _CondGradFuncGraph(util.CondBranchFuncGraph):
"""FuncGraph for the gradient function of the branch of an If op.
- Handles unwrapping optional intermediate values that are captured by the
- gradient computation.
+ Handles wrapping and unwrapping intermediate values that are captured by the
+ gradient computation in optionals.
+
+ Attributes:
+ if_op_needs_rewrite: True if any intermediates were captured, meaning the
+ forward If op needs to be written to output the wrapped intermediates.
"""
def __init__(self, name, forward_graph):
super(_CondGradFuncGraph, self).__init__(name, read_only_collections=False)
+ self.if_op_needs_rewrite = False
self._forward_graph = forward_graph
+ # Maps from forward intermediate tensor -> the unwrapped captured
+ # intermediate.
+ self._indirect_captures = {}
+ # Maps unwrapped intermediate -> optional-wrapped intermediate in the
+ # forward graph.
+ self._wrapped_intermediates = collections.OrderedDict()
+ # Raw intermediates captured from the forward graph. Populated iff we're in
+ # an XLA context.
+ self._xla_intermediates = []
+
+ @property
+ def wrapped_intermediates(self):
+ """The optional-wrapped intermediates captured from the forward graph."""
+ return list(self._wrapped_intermediates.values())
+
+ @property
+ def xla_intermediates(self):
+ """Raw intermediates captured from the forward graph if XLA is enabled."""
+ return self._xla_intermediates
def _capture_helper(self, tensor, name):
if (tensor.graph is not self._forward_graph or
@@ -554,19 +583,43 @@
tensor in self._forward_graph.outputs):
return super(_CondGradFuncGraph, self)._capture_helper(tensor, name)
- # 'tensor' is an intermediate in the forward graph. We find the corresonding
- # optional tensor, which is output from the If op, and capture it as
- # normal. We then unwrap the captured optional value to get the raw
- # intermediate value.
- for consumer in tensor.consumers():
- if (consumer.type == "OptionalFromValue"
- and consumer.outputs[0] in self._forward_graph.outputs):
- optional = consumer.outputs[0]
- captured_optional = super(_CondGradFuncGraph, self)._capture_helper(
- optional, name)
- return gen_dataset_ops.optional_get_value(
- captured_optional, [tensor.dtype], [tensor.shape])[0]
- raise ValueError(
- "Couldn't find OptionalFromValue consumer for tensor '%s'.\n"
- "This is an internal bug, please report at "
- "https://github.com/tensorflow/tensorflow/issues." % tensor.name)
+ if control_flow_util.InXlaContext(ops.get_default_graph()):
+ # XLA does not yet support optionals, so capture intermediates directly.
+ # TODO(skyewm,jpienaar): can XLA support optionals?
+ if tensor not in self.captures:
+ self.xla_intermediates.append(tensor)
+ self.if_op_needs_rewrite = True
+ return super(_CondGradFuncGraph, self)._capture_helper(tensor, name)
+
+ captured_tensor = self._indirect_captures.get(tensor)
+ if captured_tensor is not None:
+ return captured_tensor
+
+ # 'tensor' is an uncaptured intermediate in the forward graph. We wrap it in
+ # an optional in the forward graph and capture the optional normally. We
+ # then unwrap the captured optional value in the gradient graph to get the
+ # raw intermediate value.
+
+ if tensor not in self._wrapped_intermediates:
+ # If the gradient has already been computed for this If op, 'tensor' may
+ # already be wrapped.
+ for consumer in tensor.consumers():
+ if (consumer.type == "OptionalFromValue"
+ and consumer.outputs[0] in self._forward_graph.outputs):
+ optional = consumer.outputs[0]
+ break
+ else:
+ # 'tensor' hasn't been wrapped, do it now.
+ with self._forward_graph.as_default():
+ optional = gen_dataset_ops.optional_from_value([tensor])
+ self.if_op_needs_rewrite = True
+
+ self._wrapped_intermediates[tensor] = optional
+
+ optional = self._wrapped_intermediates[tensor]
+ captured_optional = super(_CondGradFuncGraph, self)._capture_helper(
+ optional, name)
+ captured_tensor = gen_dataset_ops.optional_get_value(
+ captured_optional, [tensor.dtype], [tensor.shape])[0]
+ self._indirect_captures[tensor] = captured_tensor
+ return captured_tensor
diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py
index 75c0818..99216d7 100644
--- a/tensorflow/python/ops/control_flow_ops.py
+++ b/tensorflow/python/ops/control_flow_ops.py
@@ -24,13 +24,11 @@
import abc
import collections
import functools
-import os
import six
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.protobuf import control_flow_pb2
-from tensorflow.python import tf2
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -56,6 +54,7 @@
from tensorflow.python.util import compat
from tensorflow.python.util import deprecation
from tensorflow.python.util import nest
+from tensorflow.python.util import tf_should_use
from tensorflow.python.util.lazy_loader import LazyLoader
from tensorflow.python.util.tf_export import tf_export
@@ -70,9 +69,6 @@
while_v2 = LazyLoader("while_v2", globals(),
"tensorflow.python.ops.while_v2")
-ENABLE_COND_V2 = tf2.enabled() or os.getenv("TF_ENABLE_COND_V2", "0") != "0"
-ENABLE_WHILE_V2 = tf2.enabled() or os.getenv("TF_ENABLE_WHILE_V2", "0") != "0"
-
# We override the 'tuple' for a control flow op, so we keep python's
# existing 'tuple' for later use in this module.
_basetuple = tuple
@@ -112,6 +108,7 @@
# Assert and Print are special symbols in python, so we must
# use an upper-case version of them.
@tf_export("debugging.Assert", "Assert")
+@tf_should_use.should_use_result
def Assert(condition, data, summarize=None, name=None):
"""Asserts that the given condition is true.
@@ -2050,7 +2047,7 @@
```
"""
- if ENABLE_COND_V2 and not context.executing_eagerly():
+ if util.ENABLE_CONTROL_FLOW_V2 and not context.executing_eagerly():
return cond_v2.cond_v2(pred, true_fn, false_fn, name)
# We needed to make true_fn/false_fn keyword arguments for
@@ -3485,7 +3482,7 @@
```
"""
- if ENABLE_WHILE_V2 and not context.executing_eagerly():
+ if util.ENABLE_CONTROL_FLOW_V2 and not context.executing_eagerly():
return while_v2.while_loop(
cond,
body,
diff --git a/tensorflow/python/ops/control_flow_ops_benchmark.py b/tensorflow/python/ops/control_flow_ops_benchmark.py
index 9ba5ff2..9dd1e66 100644
--- a/tensorflow/python/ops/control_flow_ops_benchmark.py
+++ b/tensorflow/python/ops/control_flow_ops_benchmark.py
@@ -27,6 +27,7 @@
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 control_flow_util
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.platform import test
@@ -94,28 +95,28 @@
iters=self.NUM_ITERS)
def benchmark_cond_v1_defun(self):
- old_val = control_flow_ops.ENABLE_COND_V2
- control_flow_ops.ENABLE_COND_V2 = False
+ old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = False
self._benchmark_defun()
- control_flow_ops.ENABLE_COND_V2 = old_val
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
def benchmark_cond_v2_defun(self):
- old_val = control_flow_ops.ENABLE_COND_V2
- control_flow_ops.ENABLE_COND_V2 = True
+ old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
self._benchmark_defun()
- control_flow_ops.ENABLE_COND_V2 = old_val
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
def benchmark_cond_v1_graph(self):
- old_val = control_flow_ops.ENABLE_COND_V2
- control_flow_ops.ENABLE_COND_V2 = False
+ old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = False
self._benchmark_graph()
- control_flow_ops.ENABLE_COND_V2 = old_val
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
def benchmark_cond_v2_graph(self):
- old_val = control_flow_ops.ENABLE_COND_V2
- control_flow_ops.ENABLE_COND_V2 = True
+ old_val = control_flow_util.ENABLE_CONTROL_FLOW_V2
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
self._benchmark_graph()
- control_flow_ops.ENABLE_COND_V2 = old_val
+ control_flow_util.ENABLE_CONTROL_FLOW_V2 = old_val
if __name__ == "__main__":
ops.enable_eager_execution()
diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py
index b19ec4b..f1dd4f5 100644
--- a/tensorflow/python/ops/control_flow_ops_test.py
+++ b/tensorflow/python/ops/control_flow_ops_test.py
@@ -21,9 +21,9 @@
import collections
import numpy as np
+from tensorflow.python import tf2
from tensorflow.core.framework import graph_pb2
from tensorflow.core.framework import node_def_pb2
-from tensorflow.python.client import session
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -127,56 +127,56 @@
node { name: "root" op: "NoOp" input: "^a" input: "^b" }
""", self._StripGraph(gd))
+ @test_util.run_deprecated_v1
def testPassingNonTensors(self):
- with ops.Graph().as_default():
- with self.assertRaises(TypeError):
- control_flow_ops.group(1, 2)
+ with self.assertRaises(TypeError):
+ control_flow_ops.group(1, 2)
class ShapeTestCase(test_util.TensorFlowTestCase):
def testShape(self):
- with ops.Graph().as_default():
- tensor = constant_op.constant([1.0, 2.0])
- self.assertEquals([2], tensor.get_shape())
- self.assertEquals([2],
- control_flow_ops.with_dependencies(
- [constant_op.constant(1.0)], tensor).get_shape())
+ tensor = constant_op.constant([1.0, 2.0])
+ self.assertEquals([2], tensor.get_shape())
+ self.assertEquals([2],
+ control_flow_ops.with_dependencies(
+ [constant_op.constant(1.0)], tensor).get_shape())
class WithDependenciesTestCase(test_util.TensorFlowTestCase):
+ @test_util.run_deprecated_v1
def testTupleDependencies(self):
- with ops.Graph().as_default():
- counter = variable_scope.get_variable(
- "my_counter", shape=[], initializer=init_ops.zeros_initializer())
- increment_counter = state_ops.assign_add(counter, 1)
- const_with_dep = control_flow_ops.with_dependencies(
- (increment_counter, constant_op.constant(42)),
- constant_op.constant(7))
- with self.cached_session():
- variables.global_variables_initializer().run()
- self.assertEquals(0, self.evaluate(counter))
- self.assertEquals(7, self.evaluate(const_with_dep))
- self.assertEquals(1, self.evaluate(counter))
+ counter = variable_scope.get_variable(
+ "my_counter", shape=[], initializer=init_ops.zeros_initializer())
+ increment_counter = state_ops.assign_add(counter, 1)
+ const_with_dep = control_flow_ops.with_dependencies(
+ (increment_counter, constant_op.constant(42)),
+ constant_op.constant(7))
+ self.evaluate(variables.global_variables_initializer())
+ self.assertEquals(0, self.evaluate(counter))
+ self.assertEquals(7, self.evaluate(const_with_dep))
+ self.assertEquals(1, self.evaluate(counter))
+
+ @test_util.run_deprecated_v1
def testListDependencies(self):
- with ops.Graph().as_default():
- counter = variable_scope.get_variable(
- "my_counter", shape=[], initializer=init_ops.zeros_initializer())
- increment_counter = state_ops.assign_add(counter, 1)
- const_with_dep = control_flow_ops.with_dependencies(
- [increment_counter, constant_op.constant(42)],
- constant_op.constant(7))
- with self.cached_session():
- variables.global_variables_initializer().run()
- self.assertEquals(0, self.evaluate(counter))
- self.assertEquals(7, self.evaluate(const_with_dep))
- self.assertEquals(1, self.evaluate(counter))
+ counter = variable_scope.get_variable(
+ "my_counter", shape=[], initializer=init_ops.zeros_initializer())
+ increment_counter = state_ops.assign_add(counter, 1)
+ const_with_dep = control_flow_ops.with_dependencies(
+ [increment_counter, constant_op.constant(42)],
+ constant_op.constant(7))
+
+ self.evaluate(variables.global_variables_initializer())
+ self.assertEquals(0, self.evaluate(counter))
+ self.assertEquals(7, self.evaluate(const_with_dep))
+ self.assertEquals(1, self.evaluate(counter))
class SwitchTestCase(test_util.TensorFlowTestCase):
+ @test_util.run_deprecated_v1
def testIndexedSlicesWithDenseShape(self):
with self.cached_session():
data = ops.IndexedSlices(
@@ -190,68 +190,64 @@
self.assertAllEqual([1, 2, 3], switch_true.values.eval())
self.assertAllEqual([0, 1], switch_true.indices.eval())
+ @test_util.run_deprecated_v1
def testIndexedSlicesGradient(self):
- with ops.Graph().as_default():
- embedding_matrix = variable_scope.get_variable(
- "embedding_matrix", [5, 5],
- initializer=init_ops.random_normal_initializer())
+ embedding_matrix = variable_scope.get_variable(
+ "embedding_matrix", [5, 5],
+ initializer=init_ops.random_normal_initializer())
- def cond(it, _):
- return it < 5
+ def cond(it, _):
+ return it < 5
- def body(it, cost):
- embedding = embedding_ops.embedding_lookup(embedding_matrix + 0.0, [0])
- cost += math_ops.reduce_sum(embedding)
- return it + 1, cost
+ def body(it, cost):
+ embedding = embedding_ops.embedding_lookup(embedding_matrix + 0.0, [0])
+ cost += math_ops.reduce_sum(embedding)
+ return it + 1, cost
- _, cost = control_flow_ops.while_loop(
- cond, body, [constant_op.constant(0),
- constant_op.constant(0.0)])
- optimizer = momentum.MomentumOptimizer(0.1, 0.9)
- train_op = optimizer.minimize(cost)
- with self.cached_session() as sess:
- self.evaluate(variables.global_variables_initializer())
- for _ in range(10):
- self.evaluate([train_op])
+ _, cost = control_flow_ops.while_loop(
+ cond, body, [constant_op.constant(0),
+ constant_op.constant(0.0)])
+ optimizer = momentum.MomentumOptimizer(0.1, 0.9)
+ train_op = optimizer.minimize(cost)
+ with self.cached_session():
+ self.evaluate(variables.global_variables_initializer())
+ for _ in range(10):
+ self.evaluate([train_op])
def testResourceReadInLoop(self):
- with ops.Graph().as_default():
- embedding_matrix = variable_scope.get_variable(
- "embedding_matrix",
- initializer=[[2.0], [3.0]],
- use_resource=True)
+ embedding_matrix = variable_scope.get_variable(
+ "embedding_matrix", initializer=[[2.0], [3.0]], use_resource=True)
- def cond(it, _):
- return it < 5
+ def cond(it, _):
+ return it < 5
- def body(it, cost):
- embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
- cost += math_ops.reduce_sum(embedding)
- return it + 1, cost
+ def body(it, cost):
+ embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
+ cost += math_ops.reduce_sum(embedding)
+ return it + 1, cost
- _, cost = control_flow_ops.while_loop(
- cond, body, [constant_op.constant(0),
- constant_op.constant(0.0)])
- with self.cached_session() as sess:
- self.evaluate(variables.global_variables_initializer())
- self.assertAllEqual(10.0, self.evaluate(cost))
+ _, cost = control_flow_ops.while_loop(
+ cond, body, [constant_op.constant(0),
+ constant_op.constant(0.0)])
+ with self.cached_session():
+ self.evaluate(variables.global_variables_initializer())
+ self.assertAllEqual(10.0, self.evaluate(cost))
def doTestIndexedSlicesGradientInCondInWhileLoop(self, use_resource=False):
- with ops.Graph().as_default():
- embedding_matrix = variable_scope.get_variable(
- "embedding_matrix", [5, 5],
- initializer=init_ops.random_normal_initializer(),
- use_resource=use_resource)
+ embedding_matrix = variable_scope.get_variable(
+ "embedding_matrix", [5, 5],
+ initializer=init_ops.random_normal_initializer(),
+ use_resource=use_resource)
- def cond(it, _):
- return it < 5
+ def cond(it, _):
+ return it < 5
- def body(it, cost):
- embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
- cost = control_flow_ops.cond(
- math_ops.equal(it, 3), lambda: math_ops.square(cost),
- lambda: cost + math_ops.reduce_sum(embedding))
- return it + 1, cost
+ def body(it, cost):
+ embedding = embedding_ops.embedding_lookup(embedding_matrix, [0])
+ cost = control_flow_ops.cond(
+ math_ops.equal(it, 3), lambda: math_ops.square(cost),
+ (lambda: cost + math_ops.reduce_sum(embedding)))
+ return it + 1, cost
_, cost = control_flow_ops.while_loop(
cond, body, [constant_op.constant(0),
@@ -269,7 +265,7 @@
static_grads = math_ops.segment_sum(static_grads.values,
static_grads.indices)
- with self.cached_session() as sess:
+ with self.cached_session():
self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(*self.evaluate([static_grads, dynamic_grads]))
@@ -279,6 +275,7 @@
def testIndexedSlicesGradientInCondInWhileLoopResource(self):
self.doTestIndexedSlicesGradientInCondInWhileLoop(use_resource=True)
+ @test_util.run_v1_only("b/120545219")
def testIndexedSlicesWithShapeGradientInWhileLoop(self):
for dtype in [dtypes.float32, dtypes.float64]:
with self.cached_session() as sess:
@@ -308,6 +305,7 @@
self.assertEquals(o, 20)
self.assertAllEqual(grad, [1] * num_steps)
+ @test_util.run_v1_only("b/120545219")
def testIndexedSlicesWithDynamicShapeGradientInWhileLoop(self):
for dtype in [dtypes.float32, dtypes.float64]:
with self.cached_session() as sess:
@@ -335,105 +333,94 @@
self.assertEquals(o, 6)
self.assertAllEqual(grad, [1] * 3)
+ @test_util.run_deprecated_v1
def testGradientThroughSingleBranchOutsideOfContext(self):
- with self.cached_session():
- x = constant_op.constant(2.)
- s = constant_op.constant(True)
- x_false, x_true = control_flow_ops.switch(x, s)
- grad_x_true = gradients_impl.gradients(x_true, x)[0]
- grad_x_false = gradients_impl.gradients(x_false, x)[0]
- self.assertEquals(grad_x_true.eval(), 1.)
- self.assertEquals(grad_x_false.eval(), 0.)
+ x = constant_op.constant(2.)
+ s = constant_op.constant(True)
+ x_false, x_true = control_flow_ops.switch(x, s)
+ grad_x_true = gradients_impl.gradients(x_true, x)[0]
+ grad_x_false = gradients_impl.gradients(x_false, x)[0]
+ self.assertEquals(self.evaluate(grad_x_true), 1.)
+ self.assertEquals(self.evaluate(grad_x_false), 0.)
class CondTest(test_util.TensorFlowTestCase):
def testCondTrue(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(2)
- y = constant_op.constant(5)
- z = control_flow_ops.cond(
- math_ops.less(x, y), lambda: math_ops.multiply(x, 17),
- lambda: math_ops.add(y, 23))
- self.assertEquals(z.eval(), 34)
+ x = constant_op.constant(2)
+ y = constant_op.constant(5)
+ z = control_flow_ops.cond(
+ math_ops.less(
+ x,
+ y), lambda: math_ops.multiply(x, 17), lambda: math_ops.add(y, 23))
+ self.assertEquals(self.evaluate(z), 34)
def testCondFalse(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(2)
- y = constant_op.constant(1)
- z = control_flow_ops.cond(
- math_ops.less(x, y), lambda: math_ops.multiply(x, 17),
- lambda: math_ops.add(y, 23))
- self.assertEquals(z.eval(), 24)
+ x = constant_op.constant(2)
+ y = constant_op.constant(1)
+ z = control_flow_ops.cond(
+ math_ops.less(
+ x,
+ y), lambda: math_ops.multiply(x, 17), lambda: math_ops.add(y, 23))
+ self.assertEquals(self.evaluate(z), 24)
def testCondTrueLegacy(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(2)
- y = constant_op.constant(5)
- z = control_flow_ops.cond(
- math_ops.less(x, y), fn1=lambda: math_ops.multiply(x, 17),
- fn2=lambda: math_ops.add(y, 23))
- self.assertEquals(z.eval(), 34)
+ x = constant_op.constant(2)
+ y = constant_op.constant(5)
+ z = control_flow_ops.cond(
+ math_ops.less(x, y),
+ fn1=lambda: math_ops.multiply(x, 17),
+ fn2=lambda: math_ops.add(y, 23))
+ self.assertEquals(self.evaluate(z), 34)
def testCondFalseLegacy(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(2)
- y = constant_op.constant(1)
- z = control_flow_ops.cond(
- math_ops.less(x, y), fn1=lambda: math_ops.multiply(x, 17),
- fn2=lambda: math_ops.add(y, 23))
- self.assertEquals(z.eval(), 24)
+ x = constant_op.constant(2)
+ y = constant_op.constant(1)
+ z = control_flow_ops.cond(
+ math_ops.less(x, y),
+ fn1=lambda: math_ops.multiply(x, 17),
+ fn2=lambda: math_ops.add(y, 23))
+ self.assertEquals(self.evaluate(z), 24)
+ @test_util.run_deprecated_v1
def testCondModifyBoolPred(self):
# This test in particular used to fail only when running in GPU, hence
# use_gpu=True.
- with ops.Graph().as_default():
- with session.Session() as sess:
- bool_var = variable_scope.get_variable("bool_var", dtype=dtypes.bool,
- initializer=True)
- cond_on_bool_var = control_flow_ops.cond(
- pred=bool_var,
- true_fn=lambda: state_ops.assign(bool_var, False),
- false_fn=lambda: True)
- self.evaluate(bool_var.initializer)
- self.assertEquals(self.evaluate(cond_on_bool_var), False)
- self.assertEquals(self.evaluate(cond_on_bool_var), True)
+ with test_util.use_gpu():
+ bool_var = variable_scope.get_variable(
+ "bool_var", dtype=dtypes.bool, initializer=True)
+ cond_on_bool_var = control_flow_ops.cond(
+ pred=bool_var,
+ true_fn=lambda: state_ops.assign(bool_var, False),
+ false_fn=lambda: True)
+ self.evaluate(bool_var.initializer)
+ self.assertEquals(self.evaluate(cond_on_bool_var), False)
+ self.assertEquals(self.evaluate(cond_on_bool_var), True)
def testCondMissingArg1(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(1)
- with self.assertRaises(TypeError):
- control_flow_ops.cond(True, false_fn=lambda: x)
+ x = constant_op.constant(1)
+ with self.assertRaises(TypeError):
+ control_flow_ops.cond(True, false_fn=lambda: x)
def testCondMissingArg2(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(1)
- with self.assertRaises(TypeError):
- control_flow_ops.cond(True, lambda: x)
+ x = constant_op.constant(1)
+ with self.assertRaises(TypeError):
+ control_flow_ops.cond(True, lambda: x)
def testCondDuplicateArg1(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(1)
- with self.assertRaises(TypeError):
- control_flow_ops.cond(True, lambda: x, lambda: x, fn1=lambda: x)
+ x = constant_op.constant(1)
+ with self.assertRaises(TypeError):
+ control_flow_ops.cond(True, lambda: x, lambda: x, fn1=lambda: x)
def testCondDuplicateArg2(self):
- with ops.Graph().as_default():
- with session.Session():
- x = constant_op.constant(1)
- with self.assertRaises(TypeError):
- control_flow_ops.cond(True, lambda: x, lambda: x, fn2=lambda: x)
+ x = constant_op.constant(1)
+ with self.assertRaises(TypeError):
+ control_flow_ops.cond(True, lambda: x, lambda: x, fn2=lambda: x)
class ContextTest(test_util.TensorFlowTestCase):
+ @test_util.run_deprecated_v1
def testCondContext(self):
with self.cached_session() as sess:
x = constant_op.constant(2)
@@ -463,12 +450,15 @@
control_flow_ops.WhileContext.from_proto(
control_flow_context.to_proto()).to_proto())
+ @test_util.run_deprecated_v1
def testWhileContext(self):
self._testWhileContextHelper()
+ @test_util.run_deprecated_v1
def testWhileContextWithMaximumIterations(self):
self._testWhileContextHelper(maximum_iterations=10)
+ @test_util.run_deprecated_v1
def testControlContextImportScope(self):
class NoABCControlFlowContext(control_flow_ops.ControlFlowContext):
"""A noop wrapper around `ControlFlowContext`.
@@ -575,7 +565,7 @@
strict=strict)
with self.cached_session() as sess:
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
true_feed_dict = {condition: True}
true_feed_dict.update(feed_dict)
result_cond, result_case = sess.run([output_cond, output_case],
@@ -591,6 +581,7 @@
if check_cond:
self.assertAllEqualNested(result_case, expected_value_false)
+ @test_util.run_deprecated_v1
def test_int(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: 1
@@ -600,6 +591,7 @@
self._testShape(fn_true, fn_false, shape, strict=True)
self._testReturnValues(fn_true, fn_false, 1, 2, strict=True)
+ @test_util.run_deprecated_v1
def test_float(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: 1.0
@@ -607,12 +599,14 @@
self._testShape(fn_true, fn_false, shape)
self._testReturnValues(fn_true, fn_false, 1.0, 2.0)
+ @test_util.run_deprecated_v1
def test_noop(self):
shape = tensor_shape.TensorShape(None)
self._testShape(control_flow_ops.no_op, control_flow_ops.no_op, shape)
self._testReturnValues(control_flow_ops.no_op, control_flow_ops.no_op,
True, False, check_cond=False)
+ @test_util.run_deprecated_v1
def test_string(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: "abc"
@@ -620,6 +614,7 @@
self._testShape(fn_true, fn_false, shape)
self._testReturnValues(fn_true, fn_false, b"abc", b"xyz")
+ @test_util.run_deprecated_v1
def test_variable(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: variables.Variable(3.0)
@@ -627,6 +622,7 @@
self._testShape(fn_true, fn_false, shape)
self._testReturnValues(fn_true, fn_false, 3.0, 4.0)
+ @test_util.run_v1_only("b/120553181")
def test_none(self):
fn_none = lambda: None
fn_tensor = lambda: constant_op.constant(1)
@@ -637,6 +633,7 @@
with self.assertRaises(ValueError):
control_flow_ops.cond(constant_op.constant(True), fn_tensor, fn_none)
+ @test_util.run_deprecated_v1
def test_tensors(self):
def _build_true_branch(dtype):
@@ -665,6 +662,7 @@
(np.zeros([2, 2]), np.ones([3, 3])),
(np.ones([2, 2]), np.zeros([3, 3])))
+ @test_util.run_deprecated_v1
def test_tensors_unknown_shape(self):
def _build_true_branch(dtype):
@@ -693,6 +691,7 @@
feed_dict={true_tensor: np.zeros([2, 2]),
false_tensor: np.ones([2, 2])})
+ @test_util.run_deprecated_v1
def test_sparse_tensors(self):
shape = tensor_shape.TensorShape([None, None])
@@ -708,11 +707,14 @@
values=[1, 2], dense_shape=[3, 4])
value2 = sparse_tensor.SparseTensorValue(indices=[[0, 0], [2, 1]],
values=[3, 4], dense_shape=[3, 4])
- self._testShape(true_fn, false_fn, shape)
- self._testReturnValues(true_fn, false_fn, value1, value2)
+ # Non-strict cond is only available in v1
+ if not tf2.enabled():
+ self._testShape(true_fn, false_fn, shape)
+ self._testReturnValues(true_fn, false_fn, value1, value2)
self._testShape(true_fn, false_fn, [shape], strict=True)
self._testReturnValues(true_fn, false_fn, [value1], [value2], strict=True)
+ @test_util.run_deprecated_v1
def test_tensors_with_partially_specified_shapes(self):
def _build_branch(dtype, shape):
@@ -742,6 +744,7 @@
true_tensors[2]: np.ones([3, 3]),
false_tensors[2]: np.ones([3, 3])})
+ @test_util.run_deprecated_v1
def test_tensor_arrays(self):
element_shape = tensor_shape.TensorShape([2])
ta1 = _create_tensor_array(4, element_shape)
@@ -751,6 +754,7 @@
fn_false = lambda: ta2
self._testShape(fn_true, fn_false, shape)
+ @test_util.run_deprecated_v1
def test_tensor_array_reads(self):
shape = tensor_shape.TensorShape([2])
ta = _create_tensor_array(4, shape)
@@ -758,6 +762,7 @@
fn_false = lambda: ta.read(1)
self._testShape(fn_true, fn_false, shape)
+ @test_util.run_deprecated_v1
def test_list(self):
shape = [tensor_shape.TensorShape([]), tensor_shape.TensorShape([]),
tensor_shape.TensorShape([])]
@@ -766,6 +771,7 @@
self._testShape(fn_true, fn_false, shape)
self._testReturnValues(fn_true, fn_false, [1, 2, 3.0], [3, 4, 5.0])
+ @test_util.run_v1_only("Non-strict cond is only available in v1")
def test_non_strict(self):
shape = tensor_shape.TensorShape([])
fn_tensor = lambda: constant_op.constant(1)
@@ -778,6 +784,7 @@
self._testReturnValues(fn_tensor, fn_tuple, 1, 3)
self._testReturnValues(fn_list, fn_tuple, 2, 3)
+ @test_util.run_v1_only("b/120553181")
def test_singleton_strict(self):
fn_tensor = lambda: constant_op.constant(1)
fn_list = lambda: [constant_op.constant(2)]
@@ -799,36 +806,46 @@
control_flow_ops.case([(constant_op.constant(True), fn_list)], fn_tuple,
strict=True)
+ @test_util.run_deprecated_v1
def test_singleton_list(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: [constant_op.constant(1)]
fn_false = lambda: [constant_op.constant(3)]
- self._testShape(fn_true, fn_false, shape)
- self._testReturnValues(fn_true, fn_false, 1, 3)
+ # Non-strict cond is only available in v1
+ if not tf2.enabled():
+ self._testShape(fn_true, fn_false, shape)
+ self._testReturnValues(fn_true, fn_false, 1, 3)
self._testShape(fn_true, fn_false, [shape], strict=True)
self._testReturnValues(fn_true, fn_false, [1], [3], strict=True)
+ @test_util.run_deprecated_v1
def test_singleton_tuple(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: (constant_op.constant(1),)
fn_false = lambda: (constant_op.constant(3),)
- self._testShape(fn_true, fn_false, shape)
- self._testReturnValues(fn_true, fn_false, 1, 3)
+ # Non-strict cond is only available in v1
+ if not tf2.enabled():
+ self._testShape(fn_true, fn_false, shape)
+ self._testReturnValues(fn_true, fn_false, 1, 3)
self._testShape(fn_true, fn_false, (shape,), strict=True)
self._testReturnValues(fn_true, fn_false, (1,), (3,),
strict=True)
+ @test_util.run_deprecated_v1
def test_singleton_namedtuple(self):
shape = tensor_shape.TensorShape([])
fn_true = lambda: SingletonTestTuple(constant_op.constant(1))
fn_false = lambda: SingletonTestTuple(constant_op.constant(3))
- self._testShape(fn_true, fn_false, shape)
- self._testReturnValues(fn_true, fn_false, 1, 3)
+ # Non-strict cond is only available in v1
+ if not tf2.enabled():
+ self._testShape(fn_true, fn_false, shape)
+ self._testReturnValues(fn_true, fn_false, 1, 3)
self._testShape(fn_true, fn_false, SingletonTestTuple(shape),
strict=True)
self._testReturnValues(fn_true, fn_false, SingletonTestTuple(1),
SingletonTestTuple(3), strict=True)
+ @test_util.run_deprecated_v1
def test_tuple(self):
shape = (tensor_shape.TensorShape([]), tensor_shape.TensorShape([]))
fn_true = lambda: (constant_op.constant(1), 2)
@@ -836,6 +853,7 @@
self._testShape(fn_true, fn_false, shape)
self._testReturnValues(fn_true, fn_false, (1, 2), (3, 4))
+ @test_util.run_deprecated_v1
def test_namedtuple(self):
shape = TestTuple(tensor_shape.TensorShape([]),
tensor_shape.TensorShape([]))
@@ -844,6 +862,7 @@
self._testShape(fn_true, fn_false, shape)
self._testReturnValues(fn_true, fn_false, TestTuple(1, 2), TestTuple(3, 4))
+ @test_util.run_deprecated_v1
def test_nested(self):
shape = [tensor_shape.TensorShape([]),
TestTuple(tensor_shape.TensorShape([]),
@@ -869,6 +888,7 @@
[11, TestTuple(12, [13, 14]),
np.ones([5, 5]), 16])
+ @test_util.run_deprecated_v1
def test_cond_inside_while_loop(self):
def body(i, matrix):
@@ -890,6 +910,7 @@
class CaseTest(test_util.TensorFlowTestCase):
+ @test_util.run_deprecated_v1
def testCase_withDefault(self):
x = array_ops.placeholder(dtype=dtypes.int32, shape=[])
conditions = [(math_ops.equal(x, 1), lambda: constant_op.constant(2)),
@@ -901,6 +922,7 @@
self.assertEqual(sess.run(output, feed_dict={x: 2}), 4)
self.assertEqual(sess.run(output, feed_dict={x: 3}), 6)
+ @test_util.run_deprecated_v1
def testCase_multiple_matches_exclusive(self):
x = array_ops.placeholder(dtype=dtypes.int32, shape=[])
conditions = [(math_ops.equal(x, 1), lambda: constant_op.constant(2)),
@@ -914,6 +936,7 @@
with self.assertRaisesRegexp(errors.InvalidArgumentError, "Input error:"):
sess.run(output, feed_dict={x: 2})
+ @test_util.run_deprecated_v1
def testCase_multiple_matches_non_exclusive(self):
x = array_ops.placeholder(dtype=dtypes.int32, shape=[])
conditions = [(math_ops.equal(x, 1), lambda: constant_op.constant(2)),
@@ -926,6 +949,7 @@
self.assertEqual(sess.run(output, feed_dict={x: 2}), 4)
self.assertEqual(sess.run(output, feed_dict={x: 3}), 8)
+ @test_util.run_deprecated_v1
def testCase_withoutDefault(self):
x = array_ops.placeholder(dtype=dtypes.int32, shape=[])
conditions = [(math_ops.equal(x, 1), lambda: constant_op.constant(2)),
@@ -939,6 +963,7 @@
with self.assertRaisesRegexp(errors.InvalidArgumentError, "Input error:"):
sess.run(output, feed_dict={x: 4})
+ @test_util.run_deprecated_v1
def testCase_withoutDefault_oneCondition(self):
x = array_ops.placeholder(dtype=dtypes.int32, shape=[])
conditions = [(math_ops.equal(x, 1), lambda: constant_op.constant(2))]
@@ -980,6 +1005,7 @@
# Expect a tuple since that is what the body returns.
self.assertEqual(self.evaluate(r), (10,))
+ @test_util.run_deprecated_v1
def testWhileLoopSameReturnShape_False(self):
i = constant_op.constant(0)
c = lambda i, _: math_ops.less(i, 10)
@@ -1005,6 +1031,7 @@
class AssertTest(test_util.TensorFlowTestCase):
+ @test_util.run_deprecated_v1
def testAssert(self):
i = constant_op.constant(0)
c = control_flow_ops.Assert(i < 10, [i, [10], [i + 1]])
diff --git a/tensorflow/python/ops/control_flow_util.py b/tensorflow/python/ops/control_flow_util.py
index cb628f4..1747f06 100644
--- a/tensorflow/python/ops/control_flow_util.py
+++ b/tensorflow/python/ops/control_flow_util.py
@@ -23,10 +23,18 @@
from __future__ import division
from __future__ import print_function
+import os
import traceback
+from tensorflow.python import tf2
from tensorflow.python.platform import tf_logging as logging
+ENABLE_CONTROL_FLOW_V2 = (tf2.enabled() or
+ os.getenv("TF_ENABLE_CONTROL_FLOW_V2", "0") != "0" or
+ os.getenv("TF_ENABLE_COND_V2", "0") != "0" or
+ os.getenv("TF_ENABLE_WHILE_V2", "0") != "0" or
+ os.getenv("TF_ENABLE_TENSOR_ARRAY_V2", "0") != "0")
+
def IsInXLAContext(op):
try:
diff --git a/tensorflow/python/ops/ctc_ops.py b/tensorflow/python/ops/ctc_ops.py
index 3a7eb93..45286f7 100644
--- a/tensorflow/python/ops/ctc_ops.py
+++ b/tensorflow/python/ops/ctc_ops.py
@@ -24,6 +24,7 @@
from tensorflow.python.framework import function
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import functional_ops
@@ -1029,7 +1030,7 @@
for the forward backward use case.
Examples:
- scan(lambda a, e: a + e, [1.0, 2.0, 3.0], 1.0) => [2.0, 3.0, 4.0]
+ scan(lambda a, e: a + e, [1.0, 2.0, 3.0], 1.0) => [2.0, 4.0, 7.0]
Multiple accumulators:
scan(lambda a, e: (a[0] + e, a[1] * e), [1.0, 2.0, 3.0], (0.0, 1.0))
@@ -1127,4 +1128,5 @@
def _get_dim(tensor, i):
"""Get value of tensor shape[i] preferring static value if available."""
- return tensor.shape[i].value or array_ops.shape(tensor)[i]
+ return tensor_shape.dimension_value(
+ tensor.shape[i]) or array_ops.shape(tensor)[i]
diff --git a/tensorflow/python/ops/custom_gradient.py b/tensorflow/python/ops/custom_gradient.py
index 1426e88..d96601ac2 100644
--- a/tensorflow/python/ops/custom_gradient.py
+++ b/tensorflow/python/ops/custom_gradient.py
@@ -236,6 +236,10 @@
original_tensors = all_tensors
with ops.get_default_graph().gradient_override_map({"IdentityN": name}):
all_tensors = array_ops.identity_n(all_tensors)
+ # Propagate handle data for happier shape inference for resource variables.
+ for i, t in enumerate(original_tensors):
+ if t.dtype == dtypes.resource and hasattr(t, "_handle_data"):
+ all_tensors[i]._handle_data = t._handle_data # pylint: disable=protected-access
tape_lib.record_operation(
f.__name__, all_tensors, original_tensors, tape_grad_fn)
for ot, t in zip(original_tensors, all_tensors):
diff --git a/tensorflow/python/ops/data_flow_ops.py b/tensorflow/python/ops/data_flow_ops.py
index 2030332..1557bdf 100644
--- a/tensorflow/python/ops/data_flow_ops.py
+++ b/tensorflow/python/ops/data_flow_ops.py
@@ -113,8 +113,9 @@
# pylint: disable=protected-access
-@tf_export("io.QueueBase", v1=["io.QueueBase", "QueueBase"])
-@deprecation.deprecated_endpoints("QueueBase")
+@tf_export("queue.QueueBase",
+ v1=["queue.QueueBase", "io.QueueBase", "QueueBase"])
+@deprecation.deprecated_endpoints(["io.QueueBase", "QueueBase"])
class QueueBase(object):
"""Base class for queue implementations.
@@ -616,8 +617,11 @@
@tf_export(
- "io.RandomShuffleQueue", v1=["io.RandomShuffleQueue", "RandomShuffleQueue"])
-@deprecation.deprecated_endpoints("RandomShuffleQueue")
+ "queue.RandomShuffleQueue",
+ v1=["queue.RandomShuffleQueue",
+ "io.RandomShuffleQueue", "RandomShuffleQueue"])
+@deprecation.deprecated_endpoints(
+ ["io.RandomShuffleQueue", "RandomShuffleQueue"])
class RandomShuffleQueue(QueueBase):
"""A queue implementation that dequeues elements in a random order.
@@ -702,7 +706,8 @@
super(RandomShuffleQueue, self).__init__(dtypes, shapes, names, queue_ref)
-@tf_export("FIFOQueue")
+@tf_export("queue.FIFOQueue", v1=["queue.FIFOQueue", "FIFOQueue"])
+@deprecation.deprecated_endpoints("FIFOQueue")
class FIFOQueue(QueueBase):
"""A queue implementation that dequeues elements in first-in first-out order.
@@ -760,8 +765,9 @@
@tf_export(
- "io.PaddingFIFOQueue", v1=["io.PaddingFIFOQueue", "PaddingFIFOQueue"])
-@deprecation.deprecated_endpoints("PaddingFIFOQueue")
+ "queue.PaddingFIFOQueue",
+ v1=["queue.PaddingFIFOQueue", "io.PaddingFIFOQueue", "PaddingFIFOQueue"])
+@deprecation.deprecated_endpoints(["io.PaddingFIFOQueue", "PaddingFIFOQueue"])
class PaddingFIFOQueue(QueueBase):
"""A FIFOQueue that supports batching variable-sized tensors by padding.
@@ -835,8 +841,9 @@
super(PaddingFIFOQueue, self).__init__(dtypes, shapes, names, queue_ref)
-@tf_export("io.PriorityQueue", v1=["io.PriorityQueue", "PriorityQueue"])
-@deprecation.deprecated_endpoints("PriorityQueue")
+@tf_export("queue.PriorityQueue",
+ v1=["queue.PriorityQueue", "io.PriorityQueue", "PriorityQueue"])
+@deprecation.deprecated_endpoints(["io.PriorityQueue", "PriorityQueue"])
class PriorityQueue(QueueBase):
"""A queue implementation that dequeues elements in prioritized order.
diff --git a/tensorflow/python/ops/embedding_ops.py b/tensorflow/python/ops/embedding_ops.py
index b398601..d0291e2 100644
--- a/tensorflow/python/ops/embedding_ops.py
+++ b/tensorflow/python/ops/embedding_ops.py
@@ -247,7 +247,7 @@
return ret
-@tf_export("nn.embedding_lookup")
+@tf_export(v1=["nn.embedding_lookup"])
def embedding_lookup(
params,
ids,
@@ -316,7 +316,66 @@
transform_fn=None)
-@tf_export("nn.embedding_lookup_sparse")
+@tf_export("nn.embedding_lookup", v1=[])
+def embedding_lookup_v2(
+ params,
+ ids,
+ partition_strategy="mod",
+ max_norm=None,
+ name=None):
+ """Looks up `ids` in a list of embedding tensors.
+
+ This function is used to perform parallel lookups on the list of
+ tensors in `params`. It is a generalization of
+ `tf.gather`, where `params` is
+ interpreted as a partitioning of a large embedding tensor. `params` may be
+ a `PartitionedVariable` as returned by using `tf.get_variable()` with a
+ partitioner.
+
+ If `len(params) > 1`, each element `id` of `ids` is partitioned between
+ the elements of `params` according to the `partition_strategy`.
+ In all strategies, if the id space does not evenly divide the number of
+ partitions, each of the first `(max_id + 1) % len(params)` partitions will
+ be assigned one more id.
+
+ If `partition_strategy` is `"mod"`, we assign each id to partition
+ `p = id % len(params)`. For instance,
+ 13 ids are split across 5 partitions as:
+ `[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]`
+
+ If `partition_strategy` is `"div"`, we assign ids to partitions in a
+ contiguous manner. In this case, 13 ids are split across 5 partitions as:
+ `[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]`
+
+ The results of the lookup are concatenated into a dense
+ tensor. The returned tensor has shape `shape(ids) + shape(params)[1:]`.
+
+ Args:
+ params: A single tensor representing the complete embedding tensor,
+ or a list of P tensors all of same shape except for the first dimension,
+ representing sharded embedding tensors. Alternatively, a
+ `PartitionedVariable`, created by partitioning along dimension 0. Each
+ element must be appropriately sized for the given `partition_strategy`.
+ ids: A `Tensor` with type `int32` or `int64` containing the ids to be looked
+ up in `params`.
+ partition_strategy: A string specifying the partitioning strategy, relevant
+ if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
+ is `"mod"`.
+ max_norm: If not `None`, each embedding is clipped if its l2-norm is
+ larger than this value.
+ name: A name for the operation (optional).
+
+ Returns:
+ A `Tensor` with the same type as the tensors in `params`.
+
+ Raises:
+ ValueError: If `params` is empty.
+ """
+ return embedding_lookup(params, ids, partition_strategy, name,
+ max_norm=max_norm)
+
+
+@tf_export(v1=["nn.embedding_lookup_sparse"])
def embedding_lookup_sparse(params,
sp_ids,
sp_weights,
@@ -491,7 +550,85 @@
return embeddings
-@tf_export("nn.safe_embedding_lookup_sparse")
+@tf_export("nn.embedding_lookup_sparse", v1=[])
+def embedding_lookup_sparse_v2(params,
+ sp_ids,
+ sp_weights,
+ partition_strategy="mod",
+ combiner=None,
+ max_norm=None,
+ name=None):
+ return embedding_lookup_sparse_v2(
+ params, sp_ids, sp_weights, partition_strategy, name, combiner, max_norm)
+
+
+embedding_lookup_sparse_v2.__doc__ = embedding_lookup_sparse.__doc__
+
+
+@tf_export("nn.safe_embedding_lookup_sparse", v1=[])
+def safe_embedding_lookup_sparse_v2(embedding_weights,
+ sparse_ids,
+ sparse_weights=None,
+ combiner="mean",
+ default_id=None,
+ max_norm=None,
+ name=None):
+ """Lookup embedding results, accounting for invalid IDs and empty features.
+
+ The partitioned embedding in `embedding_weights` must all be the same shape
+ except for the first dimension. The first dimension is allowed to vary as the
+ vocabulary size is not necessarily a multiple of `P`. `embedding_weights`
+ may be a `PartitionedVariable` as returned by using `tf.get_variable()` with a
+ partitioner.
+
+ Invalid IDs (< 0) are pruned from input IDs and weights, as well as any IDs
+ with non-positive weight. For an entry with no features, the embedding vector
+ for `default_id` is returned, or the 0-vector if `default_id` is not supplied.
+
+ The ids and weights may be multi-dimensional. Embeddings are always aggregated
+ along the last dimension.
+
+ Note: when doing embedding lookup on `embedding_weights`, "div" partition
+ strategy will be used. Support for other partition strategy will be added
+ later.
+
+ Args:
+ embedding_weights: A list of `P` float `Tensor`s or values representing
+ partitioned embedding `Tensor`s. Alternatively, a `PartitionedVariable`
+ created by partitioning along dimension 0. The total unpartitioned shape
+ should be `[e_0, e_1, ..., e_m]`, where `e_0` represents the vocab size
+ and `e_1, ..., e_m` are the embedding dimensions.
+ sparse_ids: `SparseTensor` of shape `[d_0, d_1, ..., d_n]` containing the
+ ids. `d_0` is typically batch size.
+ sparse_weights: `SparseTensor` of same shape as `sparse_ids`, containing
+ float weights corresponding to `sparse_ids`, or `None` if all weights are
+ be assumed to be 1.0.
+ combiner: A string specifying how to combine embedding results for each
+ entry. Currently "mean", "sqrtn" and "sum" are supported, with "mean" the
+ default.
+ default_id: The id to use for an entry with no features.
+ max_norm: If not `None`, all embeddings are l2-normalized to max_norm before
+ combining.
+ name: A name for this operation (optional).
+
+ Returns:
+ Dense `Tensor` of shape `[d_0, d_1, ..., d_{n-1}, e_1, ..., e_m]`.
+
+ Raises:
+ ValueError: if `embedding_weights` is empty.
+ """
+ return safe_embedding_lookup_sparse(
+ embedding_weights,
+ sparse_ids,
+ sparse_weights=sparse_weights,
+ combiner=combiner,
+ default_id=default_id,
+ name=name,
+ partition_strategy="div",
+ max_norm=max_norm)
+
+
+@tf_export(v1=["nn.safe_embedding_lookup_sparse"])
def safe_embedding_lookup_sparse(embedding_weights,
sparse_ids,
sparse_weights=None,
diff --git a/tensorflow/python/ops/functional_ops.py b/tensorflow/python/ops/functional_ops.py
index 57542e3..df4be1d 100644
--- a/tensorflow/python/ops/functional_ops.py
+++ b/tensorflow/python/ops/functional_ops.py
@@ -143,7 +143,8 @@
lambda i, a: i < n, compute, [i, a],
parallel_iterations=parallel_iterations,
back_prop=back_prop,
- swap_memory=swap_memory)
+ swap_memory=swap_memory,
+ maximum_iterations=n)
# TODO(akshayka): Remove the in_graph_mode check once caching devices are
# supported in Eager
@@ -253,7 +254,8 @@
compute, [i, a],
parallel_iterations=parallel_iterations,
back_prop=back_prop,
- swap_memory=swap_memory)
+ swap_memory=swap_memory,
+ maximum_iterations=n)
# TODO(akshayka): Remove the in_graph_mode check once caching devices are
# supported in Eager
diff --git a/tensorflow/python/ops/gradient_checker_v2.py b/tensorflow/python/ops/gradient_checker_v2.py
index cf84484..5d473ee 100644
--- a/tensorflow/python/ops/gradient_checker_v2.py
+++ b/tensorflow/python/ops/gradient_checker_v2.py
@@ -44,32 +44,42 @@
return y
-def _to_numpy(a):
- """Converts tensors to numpy arrays.
+def _eval_indexed_slices(a):
+ """Converts IndexedSlices to IndexedSlicesValue with numpy indices/values.
- Converts Tensors and EagerTensors to numpy arrays.
When eager execution is enabled, converts IndexedSlices
- to IndexedSlicesValue with numpy indices/values
+ to IndexedSlicesValue with numpy indices/values.
+
+ Args:
+ a: any value.
+
+ Returns:
+ If a is IndexedSlices and eager execution is enabled, calls numpy() on a's
+ fields. Otherwise returns a unchanged.
+ """
+ if isinstance(a, ops.IndexedSlices) and context.executing_eagerly():
+ return ops.IndexedSlicesValue(
+ indices=[x.numpy() for x in a.indices],
+ values=[x.numpy() for x in a.values],
+ dense_shape=a.dense_shape)
+ return a
+
+
+def _to_numpy(a):
+ """Converts Tensors and EagerTensors to numpy arrays.
Args:
a: any value.
Returns:
If a is EagerTensor or Tensor, returns the evaluation of a by calling
- numpy() or run().
- If a is IndexedSlices and eager execution is enabled, calls numpy() on a's
- fields. Otherwise returns a unchanged.
+ numpy() or run(). Otherwise returns a unchanged.
"""
if isinstance(a, ops.EagerTensor):
return a.numpy()
if isinstance(a, ops.Tensor):
sess = ops.get_default_session()
return sess.run(a)
- if isinstance(a, ops.IndexedSlices) and context.executing_eagerly():
- return ops.IndexedSlicesValue(
- indices=[x.numpy() for x in a.indices],
- values=[x.numpy() for x in a.values],
- dense_shape=a.dense_shape)
return a
@@ -147,6 +157,7 @@
for col in range(y_size):
dy_data_flat[col] = 1
grad = _to_numpy(grad_fn(dy_data, *xs)[0])
+ grad = _eval_indexed_slices(grad)
dy_data_flat[col] = 0
if isinstance(grad, ops.IndexedSlicesValue):
for i, v in zip(grad.indices, grad.values):
diff --git a/tensorflow/python/ops/gradient_checker_v2_test.py b/tensorflow/python/ops/gradient_checker_v2_test.py
index ce9ff47..191b2b6 100644
--- a/tensorflow/python/ops/gradient_checker_v2_test.py
+++ b/tensorflow/python/ops/gradient_checker_v2_test.py
@@ -22,7 +22,6 @@
import numpy as np
from tensorflow.python.eager import backprop
-from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
@@ -38,13 +37,17 @@
from tensorflow.python.platform import tf_logging
+def _random_complex(shape, dtype):
+ data = np.random.random_sample(shape).astype(dtype.as_numpy_dtype)
+ if dtype.is_complex:
+ data.imag = np.random.random_sample(shape)
+ return data
+
+
@test_util.run_all_in_graph_and_eager_modes
class GradientCheckerTest(test.TestCase):
def testAddSimple(self):
- # if context.executing_eagerly():
- # return
- np.random.seed(1) # Fix seed to avoid flakiness
size = (2, 3)
x1 = constant_op.constant(2.0, shape=size, name="x1")
x2 = constant_op.constant(3.0, shape=size, name="x2")
@@ -54,7 +57,6 @@
assert error < 1e-4
def testAddCustomized(self):
- np.random.seed(3) # Fix seed to avoid flakiness
size = (2, 3)
x1 = constant_op.constant(
2.0, shape=size, dtype=dtypes.float64, name="x1")
@@ -67,7 +69,6 @@
assert error < 1e-10
def testGather(self):
- np.random.seed(4) # Fix seed to avoid flakiness
def f(params):
index_values = [1, 3]
indices = constant_op.constant(index_values, name="i")
@@ -82,7 +83,6 @@
assert error < 1e-4
def testNestedGather(self):
- np.random.seed(5) # Fix seed to avoid flakiness
def f(params):
index_values = [1, 3, 5, 6]
indices = constant_op.constant(index_values, name="i")
@@ -100,33 +100,37 @@
assert error < 1e-4
def testComplexMul(self):
- if not context.executing_eagerly():
- return
c = constant_op.constant(5 + 7j, dtype=dtypes.complex64)
def f(x):
return c * x
- x = constant_op.constant(11 - 13j, dtype=dtypes.complex64)
+ x_shape = c.shape
+ x_dtype = c.dtype
+ x = constant_op.constant(_random_complex(x_shape, x_dtype))
analytical, numerical = gradient_checker.compute_gradient(
- f, [x], delta=0.1)
+ f, [x])
correct = np.array([[5, 7], [-7, 5]])
self.assertAllEqual(correct, analytical[0])
self.assertAllClose(correct, numerical[0], rtol=1e-4)
+ x = constant_op.constant(_random_complex(x_shape, x_dtype))
self.assertLess(
gradient_checker.max_error(*gradient_checker.compute_gradient(
- f, [x], delta=0.1)), 2e-4)
+ f, [x])), 3e-4)
def testComplexConj(self):
def f(x):
return math_ops.conj(x)
- x = constant_op.constant(11 - 13j, dtype=dtypes.complex64)
+ x_shape = ()
+ x_dtype = dtypes.complex64
+ x = constant_op.constant(_random_complex(x_shape, x_dtype))
analytical, numerical = gradient_checker.compute_gradient(
- f, [x], delta=0.1)
+ f, [x])
correct = np.array([[1, 0], [0, -1]])
self.assertAllEqual(correct, analytical[0])
self.assertAllClose(correct, numerical[0], rtol=2e-5)
+ x = constant_op.constant(_random_complex(x_shape, x_dtype))
self.assertLess(
gradient_checker.max_error(*gradient_checker.compute_gradient(
- f, [x], delta=0.1)), 2e-5)
+ f, [x])), 2e-5)
def testEmptySucceeds(self):
def f(x):
@@ -140,8 +144,6 @@
self.assertEqual(error, 0)
def testEmptyFails(self):
- # if not context.executing_eagerly():
- # return
@custom_gradient.custom_gradient
def id_bad_grad(x):
y = array_ops.identity(x)
@@ -279,8 +281,6 @@
return err
def testInputGradient(self):
- # if context.executing_eagerly():
- # return
self.assertLess(self._BuildAndTestMiniMNIST(0, "input"), 1e-8)
def testHiddenWeightGradient(self):
diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py
index 8cc4d92..0a70d6e 100644
--- a/tensorflow/python/ops/gradients_impl.py
+++ b/tensorflow/python/ops/gradients_impl.py
@@ -297,8 +297,12 @@
return new_grad_ys
-def IsTrainable(tensor):
- dtype = dtypes.as_dtype(tensor.dtype)
+def IsTrainable(tensor_or_dtype):
+ if isinstance(tensor_or_dtype, ops.Tensor):
+ dtype = tensor_or_dtype.dtype
+ else:
+ dtype = tensor_or_dtype
+ dtype = dtypes.as_dtype(dtype)
return dtype.base_dtype in (dtypes.float16, dtypes.float32, dtypes.float64,
dtypes.complex64, dtypes.complex128,
dtypes.resource, dtypes.variant)
@@ -322,6 +326,10 @@
ValueError: if sizes of gradients and inputs don't match.
TypeError: if type of any gradient is not valid for its input.
"""
+ # While ops have inputs added to them during the gradient computation, so we
+ # skip the below check. See while_v2 for details.
+ if op.type == "While": return
+
if len(grads) != len(op.inputs):
raise ValueError("Num gradients %d generated for op %s do not match num "
"inputs %d" % (len(grads), op.node_def, len(op.inputs)))
diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index a9058c4..c53afef 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -158,6 +158,7 @@
grads = gradients.gradients(z, [x])
self.assertTrue(all(x is not None for x in grads))
+ @test_util.run_v1_only("b/120545219")
def testBoundaryContinue(self):
# Test that we differentiate both 'x' and 'y' correctly when x is a
# predecessor of y.
@@ -169,6 +170,7 @@
self.assertTrue(all(x is not None for x in grads))
self.assertEqual(6.0, grads[0].eval())
+ @test_util.run_v1_only("b/120545219")
def testAggregationMethodAccumulateN(self):
with self.cached_session():
x = constant(1.0)
@@ -182,6 +184,7 @@
self.assertEqual(20.0, grads[0].eval())
self.assertEqual(10.0, grads[1].eval())
+ @test_util.run_v1_only("b/120545219")
def testAggregationMethodAddN(self):
with self.cached_session():
x = constant(1.0)
@@ -193,6 +196,7 @@
self.assertEqual(20.0, grads[0].eval())
self.assertEqual(10.0, grads[1].eval())
+ @test_util.run_v1_only("b/120545219")
def testAggregationMethodTree(self):
with self.cached_session():
x = constant(1.0)
@@ -239,6 +243,7 @@
[dx, dy], feed_dict={x: [1.0], dy.indices: [0], dy.values: [2.0]})
self.assertEqual(vdx, vdy)
+ @test_util.run_v1_only("b/120545219")
def testNonDifferentiableSwitchInWhileLoop(self):
with ops.Graph().as_default():
v = array_ops.placeholder(dtypes.float32, [])
@@ -270,6 +275,7 @@
gradient = gradients.gradients(graph.as_graph_element(var), var)
self.assertIsNotNone(gradient)
+ @test_util.run_v1_only("b/120545219")
def testVariableRefGradient(self):
with ops.Graph().as_default():
init = constant_op.constant(100.0)
@@ -277,6 +283,7 @@
gradient = gradients.gradients(var._ref(), var)
self.assertIsNotNone(gradient)
+ @test_util.run_v1_only("b/120545219")
def testDependentYs(self):
with self.cached_session():
x = constant_op.constant(3.0)
@@ -292,6 +299,7 @@
g = gradients.gradients([z, z2], x)
self.assertAllClose(17502.0, g[0].eval())
+ @test_util.run_v1_only("b/120545219")
def testPartialDerivatives(self):
with self.cached_session():
x = constant_op.constant(1.)
@@ -302,6 +310,7 @@
partialg = gradients.gradients(z, [x, y], stop_gradients=[x, y])
self.assertEqual([1.0, 1.0], [g.eval() for g in partialg])
+ @test_util.run_v1_only("b/120545219")
def testStopGradients(self):
def _MakeGraph(rng, stop_gradients=()):
def _FunctionOf(xs, k=3):
@@ -606,6 +615,7 @@
class HessianVectorProductTest(test_util.TensorFlowTestCase):
+ @test_util.run_v1_only("b/120545219")
def testHessianVectorProduct(self):
# Manually compute the Hessian explicitly for a low-dimensional problem
# and check that HessianVectorProduct matches multiplication by the
@@ -634,6 +644,7 @@
class HessianTest(test_util.TensorFlowTestCase):
+ @test_util.run_v1_only("b/120545219")
def testHessian1D(self):
# Manually compute the Hessian explicitly for a low-dimensional problem
# and check that `hessian` matches. Specifically, the Hessian of
@@ -651,6 +662,7 @@
hess_actual = self.evaluate(hess)
self.assertAllClose(hess_value, hess_actual)
+ @test_util.run_v1_only("b/120545219")
def testHessian1D_multi(self):
# Test the computation of the hessian with respect to multiple tensors
m = 4
@@ -671,6 +683,7 @@
for hess_value, hess_actual in zip(hess_values, hessians_actual):
self.assertAllClose(hess_value, hess_actual)
+ @test_util.run_v1_only("b/120545219")
def testHessianInvalidDimension(self):
for shape in [(10, 10), None]:
with self.cached_session(use_gpu=True):
@@ -679,6 +692,7 @@
with self.assertRaises(ValueError):
gradients.hessians(x, x)
+ @test_util.run_v1_only("b/120545219")
def testHessian2D_square_matrix(self):
# Manually compute the Hessian explicitly for a low-dimensional problem
# and check that `hessian` matches. Specifically, the Hessian of
@@ -700,6 +714,7 @@
self.assertAllEqual((m, m, m, m), hess_actual.shape)
self.assertAllClose(hess_value, hess_actual.reshape((m * m, m * m)))
+ @test_util.run_v1_only("b/120545219")
def testHessian2D_non_square_matrix(self):
m = 3
n = 4
@@ -722,6 +737,7 @@
class IndexedSlicesToTensorTest(test_util.TensorFlowTestCase):
+ @test_util.run_v1_only("b/120545219")
def testIndexedSlicesToTensor(self):
with self.cached_session():
np_val = np.random.rand(4, 4, 4, 4).astype(np.float32)
@@ -731,6 +747,7 @@
c_dense = math_ops.multiply(c_sparse, 1.0)
self.assertAllClose(np_val, self.evaluate(c_dense))
+ @test_util.run_v1_only("b/120545219")
def testIndexedSlicesToTensorList(self):
with self.cached_session():
numpy_list = []
@@ -747,6 +764,7 @@
packed_sparse = array_ops.stack(sparse_list)
self.assertAllClose(packed_dense.eval(), self.evaluate(packed_sparse))
+ @test_util.run_v1_only("b/120545219")
def testInt64Indices(self):
with self.cached_session():
np_val = np.random.rand(4, 4, 4, 4).astype(np.float32)
@@ -759,6 +777,7 @@
c_dense = math_ops.multiply(c_sparse, 1.0)
self.assertAllClose(np_val, self.evaluate(c_dense))
+ @test_util.run_v1_only("b/120545219")
def testWarnings(self):
# TODO(gunan) Reenable after this issue is fixed:
# https://github.com/google/protobuf/issues/2812
@@ -802,6 +821,7 @@
class OnlyRealGradientsTest(test_util.TensorFlowTestCase):
+ @test_util.run_v1_only("b/120545219")
def testRealOnly(self):
x = constant_op.constant(7+3j, dtype=dtypes.complex64)
y = math_ops.square(x)
@@ -814,6 +834,7 @@
class ResourceCondTest(test_util.TensorFlowTestCase):
+ @test_util.run_v1_only("b/120545219")
def testBasic(self):
gamma = resource_variable_ops.ResourceVariable(
np.random.random((3,)),
@@ -943,6 +964,7 @@
self.assertEqual(6., math_ops.reduce_sum(dx).numpy())
self.assertEqual(8., math_ops.reduce_sum(dw).numpy())
+ @test_util.run_v1_only("b/120545219")
def testCustomGradientErrorsWithNonResourceVariables(self):
def F(x, use_resource=False):
@@ -993,6 +1015,7 @@
# Smoke test to ensure numpy inputs are accepted
F(x)
+ @test_util.run_v1_only("b/120545219")
def testRVGradientsDynamicCond(self):
with self.cached_session():
alpha = resource_variable_ops.ResourceVariable(
@@ -1004,7 +1027,7 @@
conditional, lambda: alpha * 2, lambda: alpha * 3)
g, = gradients_impl.gradients(output, alpha)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllEqual(g.eval(), [2.0])
self.assertAllEqual(g.eval(feed_dict={conditional: False}), [3.0])
diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py
index 229393c..24d049b 100644
--- a/tensorflow/python/ops/image_ops_impl.py
+++ b/tensorflow/python/ops/image_ops_impl.py
@@ -2046,9 +2046,8 @@
3-D with shape `[batch, N, 4]` describing the N bounding boxes
associated with the image.
seed: An optional `int`. Defaults to `0`.
- If either `seed` or `seed2` are set to non-zero, the random number
- generator is seeded by the given `seed`. Otherwise, it is seeded by a
- random seed.
+ If `seed` is set to non-zero, the random number generator is seeded by
+ the given `seed`. Otherwise, it is seeded by a random seed.
min_object_covered: A Tensor of type `float32`. Defaults to `0.1`.
The cropped area of the image must contain at least this
fraction of any bounding box supplied. The value of this parameter should
diff --git a/tensorflow/python/ops/linalg/linalg_impl.py b/tensorflow/python/ops/linalg/linalg_impl.py
index 2c9476a9..df2bd88 100644
--- a/tensorflow/python/ops/linalg/linalg_impl.py
+++ b/tensorflow/python/ops/linalg/linalg_impl.py
@@ -44,6 +44,7 @@
eye = linalg_ops.eye
inv = linalg_ops.matrix_inverse
logm = gen_linalg_ops.matrix_logarithm
+lu = gen_linalg_ops.lu
tf_export('linalg.logm')(logm)
lstsq = linalg_ops.matrix_solve_ls
norm = linalg_ops.norm
diff --git a/tensorflow/python/ops/linalg/linear_operator.py b/tensorflow/python/ops/linalg/linear_operator.py
index 8efafda3..6be81f4 100644
--- a/tensorflow/python/ops/linalg/linear_operator.py
+++ b/tensorflow/python/ops/linalg/linear_operator.py
@@ -381,7 +381,10 @@
`Dimension` object.
"""
# Derived classes get this "for free" once .shape is implemented.
- return self.shape[-1]
+ if self.shape.rank is None:
+ return tensor_shape.Dimension(None)
+ else:
+ return self.shape.dims[-1]
def domain_dimension_tensor(self, name="domain_dimension_tensor"):
"""Dimension (in the sense of vector spaces) of the domain of this operator.
diff --git a/tensorflow/python/ops/math_grad_test.py b/tensorflow/python/ops/math_grad_test.py
index 822f897..f415e65 100644
--- a/tensorflow/python/ops/math_grad_test.py
+++ b/tensorflow/python/ops/math_grad_test.py
@@ -33,6 +33,8 @@
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
+RAISE = execution_callbacks.ExecutionCallback.RAISE
+
class SquaredDifferenceOpTest(test.TestCase):
@@ -385,7 +387,7 @@
self.assertAllClose([-2., 0., 2.], g)
def test_zero_grad_tape(self):
- with execution_callbacks.errstate(inf_or_nan=execution_callbacks.RAISE):
+ with execution_callbacks.errstate(inf_or_nan=RAISE):
x = constant_op.constant([-1, 0., 1.])
with backprop.GradientTape() as tape:
tape.watch(x)
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index f0d8bed..e656998 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -230,6 +230,7 @@
@tf_export("math.divide", "divide")
+@dispatch.add_dispatch_support
def divide(x, y, name=None):
"""Computes Python style division of `x` by `y`."""
@@ -242,6 +243,7 @@
@tf_export("math.multiply", "multiply")
+@dispatch.add_dispatch_support
def multiply(x, y, name=None):
return gen_math_ops.mul(x, y, name)
@@ -262,6 +264,7 @@
@tf_export("math.subtract", "subtract")
+@dispatch.add_dispatch_support
def subtract(x, y, name=None):
return gen_math_ops.sub(x, y, name)
@@ -347,6 +350,7 @@
@tf_export("math.pow", "pow")
+@dispatch.add_dispatch_support
def pow(x, y, name=None): # pylint: disable=redefined-builtin
r"""Computes the power of one value to another.
@@ -375,6 +379,7 @@
# pylint: disable=redefined-builtin,redefined-outer-name
@tf_export("dtypes.complex", "complex")
+@dispatch.add_dispatch_support
def complex(real, imag, name=None):
r"""Converts two real numbers to a complex number.
@@ -418,6 +423,7 @@
@tf_export("math.real", v1=["math.real", "real"])
@deprecation.deprecated_endpoints("real")
+@dispatch.add_dispatch_support
def real(input, name=None):
r"""Returns the real part of a complex (or real) tensor.
@@ -450,6 +456,7 @@
@tf_export("math.imag", v1=["math.imag", "imag"])
@deprecation.deprecated_endpoints("imag")
+@dispatch.add_dispatch_support
def imag(input, name=None):
r"""Returns the imaginary part of a complex (or real) tensor.
@@ -481,6 +488,7 @@
@tf_export("math.angle", v1=["math.angle", "angle"])
@deprecation.deprecated_endpoints("angle")
+@dispatch.add_dispatch_support
def angle(input, name=None):
r"""Returns the element-wise argument of a complex (or real) tensor.
@@ -520,6 +528,7 @@
@tf_export("math.round", "round")
+@dispatch.add_dispatch_support
def round(x, name=None): # pylint: disable=redefined-builtin
"""Rounds the values of a tensor to the nearest integer, element-wise.
@@ -547,6 +556,7 @@
@tf_export("dtypes.cast", "cast")
+@dispatch.add_dispatch_support
def cast(x, dtype, name=None):
"""Casts a tensor to a new type.
@@ -610,6 +620,7 @@
@tf_export("dtypes.saturate_cast", "saturate_cast")
+@dispatch.add_dispatch_support
def saturate_cast(value, dtype, name=None):
"""Performs a safe saturating cast of `value` to `dtype`.
@@ -935,6 +946,7 @@
@tf_export("math.truediv", "truediv")
+@dispatch.add_dispatch_support
def truediv(x, y, name=None):
"""Divides x / y elementwise (using Python 3 division operator semantics).
@@ -992,6 +1004,7 @@
@tf_export("div_no_nan")
+@dispatch.add_dispatch_support
def div_no_nan(x, y, name=None):
"""Computes an unsafe divide which returns 0 if the y is zero.
@@ -1021,6 +1034,7 @@
# TODO(aselle): Deprecate this once all internal functionality uses
# tf.truncatediv
@tf_export("math.floordiv", v1=["math.floordiv", "floordiv"])
+@dispatch.add_dispatch_support
@deprecation.deprecated_endpoints("floordiv")
def floordiv(x, y, name=None):
"""Divides `x / y` elementwise, rounding toward the most negative integer.
@@ -1050,16 +1064,11 @@
realdiv = gen_math_ops.real_div
-tf_export("realdiv")(realdiv)
truncatediv = gen_math_ops.truncate_div
-tf_export("truncatediv")(truncatediv)
# TODO(aselle): Rename this to floordiv when we can.
floor_div = gen_math_ops.floor_div
-tf_export("floor_div")(floor_div)
truncatemod = gen_math_ops.truncate_mod
-tf_export("truncatemod")(truncatemod)
floormod = gen_math_ops.floor_mod
-tf_export("floormod", "mod")(floormod)
def _mul_dispatch(x, y, name=None):
@@ -1095,6 +1104,7 @@
@tf_export("math.logical_xor", v1=["math.logical_xor", "logical_xor"])
+@dispatch.add_dispatch_support
@deprecation.deprecated_endpoints("logical_xor")
def logical_xor(x, y, name="LogicalXor"):
"""x ^ y = (x | y) & ~(x & y)."""
@@ -1277,6 +1287,7 @@
@tf_export("math.reduce_sum", "reduce_sum", v1=[])
+@dispatch.add_dispatch_support
def reduce_sum(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the sum of elements across dimensions of a tensor.
@@ -1454,7 +1465,7 @@
return cast(
reduce_sum(
# int64 reduction happens on GPU
- to_int64(gen_math_ops.not_equal(input, zero)),
+ cast(gen_math_ops.not_equal(input, zero), dtypes.int64),
axis=axis,
keepdims=keepdims),
dtype=dtype)
@@ -1524,6 +1535,7 @@
@tf_export("math.reduce_mean", "reduce_mean", v1=[])
+@dispatch.add_dispatch_support
def reduce_mean(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the mean of elements across dimensions of a tensor.
@@ -1675,6 +1687,7 @@
@tf_export("math.reduce_prod", "reduce_prod", v1=[])
+@dispatch.add_dispatch_support
def reduce_prod(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the product of elements across dimensions of a tensor.
@@ -1796,6 +1809,7 @@
@tf_export("math.reduce_min", "reduce_min", v1=[])
+@dispatch.add_dispatch_support
def reduce_min(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the minimum of elements across dimensions of a tensor.
@@ -1874,6 +1888,7 @@
@tf_export("math.reduce_max", "reduce_max", v1=[])
+@dispatch.add_dispatch_support
def reduce_max(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the maximum of elements across dimensions of a tensor.
@@ -1961,6 +1976,7 @@
@tf_export("reduce_all", "math.reduce_all", v1=[])
+@dispatch.add_dispatch_support
def reduce_all(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the "logical and" of elements across dimensions of a tensor.
@@ -2057,6 +2073,7 @@
@tf_export("math.reduce_any", "reduce_any", v1=[])
+@dispatch.add_dispatch_support
def reduce_any(input_tensor, axis=None, keepdims=False, name=None):
"""Computes the "logical or" of elements across dimensions of a tensor.
@@ -2619,6 +2636,7 @@
@tf_export("math.add_n", "add_n")
+@dispatch.add_dispatch_support
def add_n(inputs, name=None):
"""Adds all input tensors element-wise.
@@ -2764,6 +2782,7 @@
@tf_export("math.log_sigmoid", v1=["math.log_sigmoid", "log_sigmoid"])
+@dispatch.add_dispatch_support
@deprecation.deprecated_endpoints("log_sigmoid")
def log_sigmoid(x, name=None):
"""Computes log sigmoid of `x` element-wise.
@@ -2973,6 +2992,7 @@
@tf_export("math.conj", v1=["math.conj", "conj"])
+@dispatch.add_dispatch_support
@deprecation.deprecated_endpoints("conj")
def conj(x, name=None):
r"""Returns the complex conjugate of a complex number.
@@ -3077,6 +3097,7 @@
"math.unsorted_segment_mean",
v1=["math.unsorted_segment_mean", "unsorted_segment_mean"])
@deprecation.deprecated_endpoints("unsorted_segment_mean")
+@dispatch.add_dispatch_support
def unsorted_segment_mean(data, segment_ids, num_segments, name=None):
r"""Computes the mean along segments of a tensor.
@@ -3122,6 +3143,7 @@
"math.unsorted_segment_sqrt_n",
v1=["math.unsorted_segment_sqrt_n", "unsorted_segment_sqrt_n"])
@deprecation.deprecated_endpoints("unsorted_segment_sqrt_n")
+@dispatch.add_dispatch_support
def unsorted_segment_sqrt_n(data, segment_ids, num_segments, name=None):
r"""Computes the sum along segments of a tensor divided by the sqrt(N).
diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py
index add1621..e185dbc 100644
--- a/tensorflow/python/ops/math_ops_test.py
+++ b/tensorflow/python/ops/math_ops_test.py
@@ -223,7 +223,7 @@
@test_util.run_in_graph_and_eager_modes
def testSquaredDifference(self):
- for dtype in [np.int32, np.float16]:
+ for dtype in [np.float16, np.float32, np.float64, np.int32, np.int64]:
x = np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype)
y = np.array([-3, -2, -1], dtype=dtype)
z = (x - y) * (x - y)
@@ -231,6 +231,17 @@
z_tf = self.evaluate(math_ops.squared_difference(x, y))
self.assertAllClose(z, z_tf)
+ @test_util.run_in_graph_and_eager_modes()
+ def testComplexSquaredDifference(self):
+ for dtype in [np.complex64, np.complex128]:
+ x = np.array([[1 + 3j, 2 + 2j, 3 + 1j], [4 - 1j, 5 - 2j, 6 - 3j]],
+ dtype=dtype)
+ y = np.array([-3 + 1j, -2 + 2j, -1 + 3j], dtype=dtype)
+ z = np.conj(x - y) * (x - y)
+ with test_util.device(use_gpu=False):
+ z_tf = self.evaluate(math_ops.squared_difference(x, y))
+ self.assertAllClose(z, z_tf)
+
class ApproximateEqualTest(test_util.TensorFlowTestCase):
diff --git a/tensorflow/python/ops/metrics_impl.py b/tensorflow/python/ops/metrics_impl.py
index cb42199..ec39b17 100644
--- a/tensorflow/python/ops/metrics_impl.py
+++ b/tensorflow/python/ops/metrics_impl.py
@@ -18,6 +18,7 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
@@ -34,7 +35,6 @@
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import weights_broadcast_ops
from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.util.deprecation import deprecated
from tensorflow.python.util.tf_export import tf_export
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index 8f74f83..48dcab4 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -262,7 +262,7 @@
name=name)
-@tf_export("nn.relu_layer")
+@tf_export(v1=["nn.relu_layer"])
def relu_layer(x, weights, biases, name=None):
"""Computes Relu(x * weight + biases).
@@ -1677,7 +1677,98 @@
return _sum_rows(sampled_losses)
-@tf_export("nn.sampled_softmax_loss")
+@tf_export("nn.sampled_softmax_loss", v1=[])
+def sampled_softmax_loss_v2(weights,
+ biases,
+ labels,
+ inputs,
+ num_sampled,
+ num_classes,
+ num_true=1,
+ sampled_values=None,
+ remove_accidental_hits=True,
+ seed=None,
+ name="sampled_softmax_loss"):
+ """Computes and returns the sampled softmax training loss.
+
+ This is a faster way to train a softmax classifier over a huge number of
+ classes.
+
+ This operation is for training only. It is generally an underestimate of
+ the full softmax loss.
+
+ A common use case is to use this method for training, and calculate the full
+ sigmoid loss for evaluation or inference as in the following example:
+
+ ```python
+ if mode == "train":
+ loss = tf.nn.sampled_softmax_loss(
+ weights=weights,
+ biases=biases,
+ labels=labels,
+ inputs=inputs,
+ ...)
+ elif mode == "eval":
+ logits = tf.matmul(inputs, tf.transpose(weights))
+ logits = tf.nn.bias_add(logits, biases)
+ labels_one_hot = tf.one_hot(labels, n_classes)
+ loss = tf.nn.softmax_cross_entropy_with_logits_v2(
+ labels=labels_one_hot,
+ logits=logits)
+ ```
+
+ See our [Candidate Sampling Algorithms Reference]
+ (https://www.tensorflow.org/extras/candidate_sampling.pdf)
+
+ Also see Section 3 of [Jean et al., 2014](http://arxiv.org/abs/1412.2007)
+ ([pdf](http://arxiv.org/pdf/1412.2007.pdf)) for the math.
+
+ Note: when doing embedding lookup on `weights` and `bias`, "div" partition
+ strategy will be used. Support for other partition strategy will be added
+ later.
+
+ Args:
+ weights: A `Tensor` of shape `[num_classes, dim]`, or a list of `Tensor`
+ objects whose concatenation along dimension 0 has shape [num_classes,
+ dim]. The (possibly-sharded) class embeddings.
+ biases: A `Tensor` of shape `[num_classes]`. The class biases.
+ labels: A `Tensor` of type `int64` and shape `[batch_size, num_true]`. The
+ target classes. Note that this format differs from the `labels` argument
+ of `nn.softmax_cross_entropy_with_logits_v2`.
+ inputs: A `Tensor` of shape `[batch_size, dim]`. The forward activations of
+ the input network.
+ num_sampled: An `int`. The number of classes to randomly sample per batch.
+ num_classes: An `int`. The number of possible classes.
+ num_true: An `int`. The number of target classes per training example.
+ sampled_values: a tuple of (`sampled_candidates`, `true_expected_count`,
+ `sampled_expected_count`) returned by a `*_candidate_sampler` function.
+ (if None, we default to `log_uniform_candidate_sampler`)
+ remove_accidental_hits: A `bool`. whether to remove "accidental hits"
+ where a sampled class equals one of the target classes. Default is True.
+ seed: random seed for candidate sampling. Default to None, which doesn't set
+ the op-level random seed for candidate sampling.
+ name: A name for the operation (optional).
+
+ Returns:
+ A `batch_size` 1-D tensor of per-example sampled softmax losses.
+
+ """
+ return sampled_softmax_loss(
+ weights,
+ biases,
+ labels,
+ inputs,
+ num_sampled,
+ num_classes,
+ num_true=num_true,
+ sampled_values=sampled_values,
+ remove_accidental_hits=remove_accidental_hits,
+ partition_strategy="div",
+ name=name,
+ seed=seed)
+
+
+@tf_export(v1=["nn.sampled_softmax_loss"])
def sampled_softmax_loss(weights,
biases,
labels,
diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 472d614..611bfda 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -3040,7 +3040,7 @@
noise_shape, seed=seed, dtype=x.dtype)
# 0. if [keep_prob, 1.0) and 1. if [1.0, 1.0 + keep_prob)
binary_tensor = math_ops.floor(random_tensor)
- ret = math_ops.div(x, keep_prob) * binary_tensor
+ ret = math_ops.divide(x, keep_prob) * binary_tensor
if not context.executing_eagerly():
ret.set_shape(x.get_shape())
return ret
@@ -3776,7 +3776,7 @@
name=name))
-@tf_export("math.in_top_k", "nn.in_top_k")
+@tf_export(v1=["math.in_top_k", "nn.in_top_k"])
def in_top_k(predictions, targets, k, name=None):
r"""Says whether the targets are in the top `K` predictions.
@@ -3810,6 +3810,14 @@
return gen_nn_ops.in_top_kv2(predictions, targets, k, name=name)
+@tf_export("math.in_top_k", "nn.in_top_k", v1=[])
+def in_top_k_v2(targets, predictions, k, name=None):
+ return in_top_k(predictions, targets, k, name)
+
+
+in_top_k_v2.__doc__ = in_top_k.__doc__
+
+
tf_export(v1=["nn.quantized_avg_pool"])(gen_nn_ops.quantized_avg_pool)
tf_export(v1=["nn.quantized_conv2d"])(gen_nn_ops.quantized_conv2d)
tf_export(v1=["nn.quantized_relu_x"])(gen_nn_ops.quantized_relu_x)
diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index 0336d0d..82fab74 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -450,6 +450,7 @@
with self.assertRaises(ValueError):
nn_ops.dropout(t, array_ops.placeholder(dtypes.float32, shape=[2]))
+ @test_util.run_deprecated_v1
def testInvalidRate(self):
x_dim = 40
y_dim = 30
@@ -856,7 +857,7 @@
exp_sampled_softmax_loss = _SoftmaxCrossEntropyWithLogits(
exp_logits, exp_labels)
- got_sampled_softmax_loss = nn_impl.sampled_softmax_loss(
+ got_sampled_softmax_loss = nn_impl.sampled_softmax_loss_v2(
weights=constant_op.constant(weights),
biases=constant_op.constant(biases),
labels=constant_op.constant(labels, shape=(batch_size, 1)),
@@ -865,8 +866,7 @@
num_classes=num_classes,
num_true=1,
sampled_values=sampled_vals,
- remove_accidental_hits=False,
- partition_strategy="div")
+ remove_accidental_hits=False)
self.assertAllClose(exp_sampled_softmax_loss,
self.evaluate(got_sampled_softmax_loss), 1e-4)
@@ -874,7 +874,7 @@
# Test with sharded weights and sharded biases.
weight_shards, bias_shards = self._ShardTestEmbeddings(
weights, biases, num_shards=3)
- got_sampled_softmax_loss = nn_impl.sampled_softmax_loss(
+ got_sampled_softmax_loss = nn_impl.sampled_softmax_loss_v2(
weights=[constant_op.constant(shard) for shard in weight_shards],
biases=[constant_op.constant(shard) for shard in bias_shards],
labels=constant_op.constant(labels, shape=(batch_size, 1)),
@@ -883,8 +883,7 @@
num_classes=num_classes,
num_true=1,
sampled_values=sampled_vals,
- remove_accidental_hits=False,
- partition_strategy="div")
+ remove_accidental_hits=False)
self.assertAllClose(exp_sampled_softmax_loss,
self.evaluate(got_sampled_softmax_loss), 1e-4)
@@ -925,7 +924,7 @@
sampled_vals_bf16 = (sampled, true_exp_bf16, sampled_exp_bf16)
got_sampled_softmax_loss = math_ops.cast(
- nn_impl.sampled_softmax_loss(
+ nn_impl.sampled_softmax_loss_v2(
weights=constant_op.constant(weights, dtype=dtypes.bfloat16),
biases=constant_op.constant(biases, dtype=dtypes.bfloat16),
labels=constant_op.constant(
@@ -935,8 +934,7 @@
num_classes=num_classes,
num_true=1,
sampled_values=sampled_vals_bf16,
- remove_accidental_hits=False,
- partition_strategy="div"), dtypes.float32)
+ remove_accidental_hits=False), dtypes.float32)
self.assertAllClose(exp_sampled_softmax_loss,
self.evaluate(got_sampled_softmax_loss), 1e-1)
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 cc20d7c..933bddd 100644
--- a/tensorflow/python/ops/parallel_for/control_flow_ops_test.py
+++ b/tensorflow/python/ops/parallel_for/control_flow_ops_test.py
@@ -927,7 +927,10 @@
outputs[1] = constant_op.constant(0.)
outputs[2] = constant_op.constant(0.)
loss = nn.l2_loss(outputs[0])
- gradients = g.gradient(loss, [x1, scale, offset])
+ if is_training:
+ gradients = g.gradient(loss, [x1, scale, offset])
+ else:
+ gradients = [constant_op.constant(0.)] * 3
return outputs + gradients
# pylint: enable=cell-var-from-loop
diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index e335c5c..89b8c4a 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -1,3 +1,5 @@
+load("//tensorflow:tensorflow.bzl", "py_test")
+
package(
default_visibility = [
"//intelligence/datum/prensor:__pkg__",
@@ -11,8 +13,6 @@
exports_files(["LICENSE"])
-load("//tensorflow:tensorflow.bzl", "py_test")
-
#-------------------------------------------------------------------------------
# RaggedTensor
#-------------------------------------------------------------------------------
@@ -25,7 +25,7 @@
deps = [
":ragged_array_ops",
":ragged_conversion_ops",
- ":ragged_elementwise_ops",
+ ":ragged_dispatch",
":ragged_factory_ops",
":ragged_functional_ops",
":ragged_getitem",
@@ -48,7 +48,6 @@
srcs_version = "PY2AND3",
deps = [
":ragged_conversion_ops",
- ":ragged_factory_ops",
":ragged_functional_ops",
":ragged_math_ops",
":ragged_tensor",
@@ -82,6 +81,7 @@
"//tensorflow/python:ragged_conversion_ops_gen",
"//tensorflow/python:sparse_tensor",
"//tensorflow/python:tensor_shape",
+ "//tensorflow/python:util",
],
)
@@ -95,6 +95,7 @@
"//tensorflow/python:framework_ops",
"//tensorflow/python:math_ops",
"//tensorflow/python:tensor_util",
+ "//tensorflow/python:util",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/ops/ragged:ragged_tensor_value",
"//third_party/py/numpy",
@@ -110,6 +111,7 @@
":ragged_tensor",
":ragged_util",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python:util",
],
)
@@ -147,24 +149,6 @@
"//tensorflow/python:math_ops",
"//tensorflow/python:ragged_math_ops_gen",
"//tensorflow/python:tensor_util",
- ],
-)
-
-py_library(
- name = "ragged_elementwise_ops",
- srcs = ["ragged_elementwise_ops.py"],
- srcs_version = "PY2AND3",
- deps = [
- ":ragged_factory_ops",
- ":ragged_tensor",
- ":ragged_tensor_shape",
- ":ragged_util",
- "//tensorflow/python:array_ops",
- "//tensorflow/python:clip_ops",
- "//tensorflow/python:framework_ops",
- "//tensorflow/python:math_ops",
- "//tensorflow/python:parsing_ops",
- "//tensorflow/python:string_ops",
"//tensorflow/python:util",
],
)
@@ -174,9 +158,9 @@
srcs = ["ragged_operators.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged_elementwise_ops",
":ragged_getitem",
":ragged_tensor",
+ "//tensorflow/python:math_ops",
"//tensorflow/python:util",
],
)
@@ -191,7 +175,11 @@
":ragged_factory_ops",
":ragged_tensor",
"//tensorflow/python:array_ops",
+ "//tensorflow/python:dtypes",
+ "//tensorflow/python:framework_ops",
"//tensorflow/python:math_ops",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:string_ops",
"//tensorflow/python:util",
],
)
@@ -202,9 +190,13 @@
srcs_version = "PY2AND3",
deps = [
":ragged_tensor_value",
+ ":ragged_util",
+ ":segment_id_ops",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python:ragged_conversion_ops_gen",
"//tensorflow/python:session",
"//tensorflow/python:tensor_shape",
+ "//tensorflow/python:util",
],
)
@@ -219,10 +211,11 @@
":ragged_tensor",
":ragged_util",
"//tensorflow/python:array_ops",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:control_flow_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_ops",
"//tensorflow/python:math_ops",
- "//tensorflow/python:tensor_shape",
"//tensorflow/python:tensor_util",
],
)
@@ -231,7 +224,10 @@
name = "ragged_tensor_value",
srcs = ["ragged_tensor_value.py"],
srcs_version = "PY2AND3",
- deps = ["//third_party/py/numpy"],
+ deps = [
+ "//tensorflow/python:util",
+ "//third_party/py/numpy",
+ ],
)
py_library(
@@ -260,6 +256,7 @@
"//tensorflow/python:math_ops",
"//tensorflow/python:tensor_shape",
"//tensorflow/python:tensor_util",
+ "//tensorflow/python:util",
],
)
@@ -268,27 +265,63 @@
srcs = ["ragged_map_ops.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged_array_ops",
- ":ragged_factory_ops",
":ragged_tensor",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:control_flow_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
"//tensorflow/python:platform",
"//tensorflow/python:sparse_tensor",
"//tensorflow/python:tensor_array_ops",
+ "//tensorflow/python:tensor_shape",
"//tensorflow/python:util",
"//tensorflow/python:variable_scope",
"//tensorflow/python/eager:context",
],
)
+py_library(
+ name = "ragged_dispatch",
+ srcs = ["ragged_dispatch.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":ragged_array_ops",
+ ":ragged_math_ops",
+ ":ragged_tensor",
+ ":ragged_tensor_shape",
+ ":ragged_util",
+ "//tensorflow/python:array_ops",
+ "//tensorflow/python:clip_ops",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:parsing_ops",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python:util",
+ "//tensorflow/python:variables",
+ "//third_party/py/numpy",
+ ],
+)
+
#-------------------------------------------------------------------------------
# RaggedTensor Tests
#-------------------------------------------------------------------------------
+py_library(
+ name = "ragged_test_util",
+ srcs = ["ragged_test_util.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":ragged_tensor",
+ ":ragged_tensor_value",
+ "//tensorflow/python:framework_ops",
+ "//tensorflow/python:framework_test_lib",
+ "//third_party/py/numpy",
+ ],
+)
+
py_test(
name = "ragged_tensor_test",
size = "medium",
@@ -298,13 +331,20 @@
"no_windows",
],
deps = [
- ":ragged",
+ ":ragged", # fixdeps: keep
+ ":ragged_factory_ops",
+ ":ragged_math_ops",
+ ":ragged_tensor",
+ ":ragged_tensor_value",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
"//tensorflow/python:framework_ops",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
@@ -316,9 +356,9 @@
srcs = ["ragged_eager_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:framework_ops",
- "//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
"@absl_py//absl/testing:parameterized",
],
@@ -329,7 +369,8 @@
srcs = ["ragged_range_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_math_ops",
+ ":ragged_test_util",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
@@ -341,7 +382,9 @@
srcs = ["ragged_tensor_bounding_shape_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
],
@@ -352,7 +395,10 @@
srcs = ["ragged_row_lengths_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
+ "//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
"@absl_py//absl/testing:parameterized",
@@ -364,7 +410,9 @@
srcs = ["ragged_gather_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
@@ -372,6 +420,7 @@
"//tensorflow/python:framework_ops",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
],
)
@@ -380,12 +429,16 @@
srcs = ["ragged_batch_gather_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
"@absl_py//absl/testing:parameterized",
],
)
@@ -395,11 +448,15 @@
srcs = ["ragged_gather_nd_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:dtypes",
+ "//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
@@ -410,7 +467,8 @@
srcs = ["ragged_row_splits_to_segment_ids_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_test_util",
+ ":segment_id_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
@@ -422,7 +480,8 @@
srcs = ["ragged_segment_ids_to_row_splits_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_test_util",
+ ":segment_id_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
@@ -434,7 +493,8 @@
srcs = ["ragged_from_tensor_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:framework_test_lib",
@@ -452,13 +512,19 @@
"no_windows",
],
deps = [
- ":ragged",
+ ":ragged", # fixdeps: keep
+ ":ragged_factory_ops",
+ ":ragged_functional_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:gradients_impl",
+ "//tensorflow/python:math_ops",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
],
)
@@ -467,13 +533,15 @@
srcs = ["ragged_from_sparse_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
"//tensorflow/python:sparse_tensor",
+ "//tensorflow/python/eager:context",
],
)
@@ -482,7 +550,8 @@
srcs = ["ragged_to_tensor_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:framework_ops",
"//tensorflow/python:framework_test_lib",
@@ -496,7 +565,10 @@
srcs = ["ragged_segment_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_math_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:errors",
@@ -511,23 +583,29 @@
srcs = ["ragged_reduce_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_math_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
- name = "ragged_map_inner_values_op_test",
- srcs = ["ragged_map_inner_values_op_test.py"],
+ name = "ragged_map_flat_values_op_test",
+ srcs = ["ragged_map_flat_values_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_functional_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
@@ -535,7 +613,6 @@
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:math_ops",
"//tensorflow/python:platform_test",
- "@absl_py//absl/testing:parameterized",
],
)
@@ -544,8 +621,10 @@
srcs = ["ragged_const_op_test.py"],
srcs_version = "PY2AND3",
deps = [
+ ":ragged",
":ragged_factory_ops",
":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
@@ -561,7 +640,10 @@
"no_windows",
],
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_tensor_value",
+ ":ragged_test_util",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
"//third_party/py/numpy",
@@ -574,7 +656,9 @@
srcs = ["convert_to_tensor_or_ragged_tensor_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_test_lib",
@@ -589,12 +673,15 @@
srcs = ["ragged_boolean_mask_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
"@absl_py//absl/testing:parameterized",
],
)
@@ -604,13 +691,16 @@
srcs = ["ragged_concat_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//tensorflow/python/eager:context",
"@absl_py//absl/testing:parameterized",
],
)
@@ -620,7 +710,9 @@
srcs = ["ragged_stack_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:constant_op",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
@@ -635,6 +727,7 @@
deps = [
":ragged_array_ops",
":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
@@ -649,6 +742,7 @@
srcs = ["ragged_util_test.py"],
srcs_version = "PY2AND3",
deps = [
+ ":ragged_test_util",
":ragged_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:constant_op",
@@ -664,7 +758,9 @@
srcs = ["ragged_expand_dims_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
"@absl_py//absl/testing:parameterized",
@@ -676,7 +772,9 @@
srcs = ["ragged_where_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged_array_ops",
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
"@absl_py//absl/testing:parameterized",
@@ -684,17 +782,26 @@
)
py_test(
- name = "ragged_elementwise_ops_test",
- srcs = ["ragged_elementwise_ops_test.py"],
+ name = "ragged_dispatch_test",
+ srcs = ["ragged_dispatch_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged", # fixdeps: keep
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
+ "//tensorflow/python:clip_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:errors",
"//tensorflow/python:framework_ops",
"//tensorflow/python:framework_test_lib",
+ "//tensorflow/python:math_ops",
+ "//tensorflow/python:parsing_ops",
"//tensorflow/python:platform_test",
+ "//tensorflow/python:sparse_tensor",
+ "//tensorflow/python:string_ops",
+ "//tensorflow/python/eager:context",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
@@ -705,7 +812,9 @@
srcs = ["ragged_operators_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged", # fixdeps: keep
+ ":ragged_factory_ops",
+ ":ragged_test_util",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
],
@@ -717,7 +826,13 @@
srcs = ["ragged_map_fn_op_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged", # fixdeps: keep
+ ":ragged_factory_ops",
+ ":ragged_functional_ops",
+ ":ragged_map_ops",
+ ":ragged_math_ops",
+ ":ragged_tensor",
+ ":ragged_test_util",
"//tensorflow/python:array_ops",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_test_lib",
@@ -725,6 +840,7 @@
"//tensorflow/python:platform_test",
"//tensorflow/python:string_ops",
"//tensorflow/python/keras:backend",
+ "//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
@@ -734,9 +850,15 @@
srcs = ["ragged_tensor_shape_test.py"],
srcs_version = "PY2AND3",
deps = [
- ":ragged",
+ ":ragged", # fixdeps: keep
+ ":ragged_factory_ops",
+ ":ragged_tensor",
+ ":ragged_tensor_shape",
+ ":ragged_test_util",
+ "//tensorflow/python:dtypes",
"//tensorflow/python:framework_test_lib",
"//tensorflow/python:platform_test",
+ "//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
diff --git a/tensorflow/python/ops/ragged/__init__.py b/tensorflow/python/ops/ragged/__init__.py
index 1b2a7be..7806f56 100644
--- a/tensorflow/python/ops/ragged/__init__.py
+++ b/tensorflow/python/ops/ragged/__init__.py
@@ -1,236 +1,47 @@
+# 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.
+# ==============================================================================
"""Ragged Tensors.
-This package defines the [`RaggedTensor`](ragged/RaggedTensor.md) class, which
-represents tensors with non-uniform shapes. In particular, each `RaggedTensor`
+This package defines ops for manipulating ragged tensors (`tf.RaggedTensor`),
+which are tensors with non-uniform shapes. In particular, each `RaggedTensor`
has one or more *ragged dimensions*, which are dimensions whose slices may have
different lengths. For example, the inner (column) dimension of
`rt=[[3, 1, 4, 1], [], [5, 9, 2], [6], []]` is ragged, since the column slices
(`rt[0, :]`, ..., `rt[4, :]`) have different lengths. For a more detailed
-description of ragged tensors, see the [`RaggedTensor`](ragged/RaggedTensor.md)
-class documentation.
-
-## RaggedTensor Operations
-
-This package also defines a collection of operations for manipulating
-ragged tensors.
-
-### RaggedTensor Versions of Standard Tensor Operations
-
-Many of the operations defined by this package are analogous to
-[`Tensor`](https://www.tensorflow.org/api_docs/python/tf/Tensor)
-operations, but they accept `RaggedTensor`s as input and can return
-`RaggedTensor`s as output. For example, `ragged.add` performs elementwise
-addition just like `tf.add`, but can be used on `RaggedTensor`s.
-
-These `RaggedTensor` versions of the standard `Tensor` operations can also be
-used with standard `Tensors`; and for the most part, they will return the same
-value that the standard `Tensor` operation would return. However, there are
-a few notable exceptions:
-
-* For [`ragged.stack(...)`](ragged/stack.md) and
- [`ragged.concat(...)`](ragged/concat.md), the input tensors are not required
- to have matching shapes. In the returned tensor, all dimensions up to the
- `axis` dimension will be ragged.
-
-### Ragged-Tensor Specific Operations
-
-The following operations are specific to ragged tensors:
-
-* **Factory ops**:
- [`constant(...)`](ragged/constant.md),
- [`from_row_splits(...)`](ragged/from_row_splits.md),
- [`from_row_lengths(...)`](ragged/from_row_lengths.md),
- [`from_row_starts(...)`](ragged/from_row_starts.md),
- [`from_row_limits(...)`](ragged/from_row_limits.md),
- [`from_value_rowids(...)`](ragged/from_value_rowids.md),
- [`from_nested_row_splits(...)`](ragged/from_nested_row_splits.md),
- [`from_nested_value_rowids(...)`](ragged/from_nested_value_rowids.md).
-
-* **Conversion ops**:
- [`from_tensor(...)`](ragged/from_tensor.md),
- [`to_tensor(...)`](ragged/to_tensor.md),
- [`from_sparse(...)`](ragged/from_sparse.md),
- [`to_sparse(...)`](ragged/to_sparse.md),
- [`from_variant(...)`](ragged/from_variant.md),
- [`to_variant(...)`](ragged/to_variant.md),
- [`convert_to_tensor_or_ragged_tensor(...)`](
- ragged/convert_to_tensor_or_ragged_tensor.md).
-
-* **Shape ops**:
- [`row_splits(...)`](ragged/row_splits.md),
- [`row_lengths(...)`](ragged/row_lengths.md),
- [`row_starts(...)`](ragged/row_starts.md),
- [`row_limits(...)`](ragged/row_limits.md),
- [`value_rowids(...)`](ragged/value_rowids.md),
- [`nrows(...)`](ragged/nrows.md),
- [`nested_row_splits(...)`](ragged/nested_row_splits.md),
- [`row_splits_to_segment_ids(...)`](ragged/row_splits_to_segment_ids.md),
- [`segment_ids_to_row_splits(...)`](ragged/segment_ids_to_row_splits.md),
- [`bounding_shape(...)`](ragged/bounding_shape.md).
-
-* **Functional ops**:
- [`map_inner_values(...)`](ragged/map_inner_values.md),
- [`make_elementwise_op(...)`](ragged/make_elementwise_op.md).
-
-
-<!-- Ragged Classes & related helper functions -->
-@@RaggedTensor
-@@RaggedTensorType
-@@RaggedTensorValue
-@@is_ragged
-
-<!-- Factory Ops -->
-@@constant
-@@constant_value
-@@from_row_splits
-@@from_row_lengths
-@@from_row_starts
-@@from_row_limits
-@@from_value_rowids
-@@from_nested_row_splits
-@@from_nested_value_rowids
-@@convert_to_tensor_or_ragged_tensor
-
-<!-- Conversion Ops -->
-@@from_tensor
-@@to_tensor
-@@from_sparse
-@@to_sparse
-@@row_splits_to_segment_ids
-@@segment_ids_to_row_splits
-
-<!-- Array Ops -->
-@@row_splits
-@@row_lengths
-@@row_starts
-@@row_limits
-@@value_rowids
-@@nrows
-@@nested_row_splits
-@@bounding_shape
-@@gather
-@@batch_gather
-@@gather_nd
-@@boolean_mask
-@@concat
-@@stack
-@@tile
-@@expand_dims
-@@where
-
-<!-- Math Ops -->
-@@range
-
-@@segment_sum
-@@segment_prod
-@@segment_min
-@@segment_max
-@@segment_mean
-@@segment_sqrt_n
-
-@@reduce_sum
-@@reduce_prod
-@@reduce_min
-@@reduce_max
-@@reduce_mean
-@@reduce_all
-@@reduce_any
-
-<!-- Functional Ops -->
-@@map_inner_values
-@@map_fn
-
-<!-- Elementwise Ops -->
-@@make_elementwise_op
-
-<!-- Shape & broadcasting -->
-@@RaggedTensorDynamicShape
-@@broadcast_to
-@@broadcast_dynamic_shape
-
-<!-- Symbols from ragged_elementwise_ops._symbols_to_export are whitelisted -->
+description of ragged tensors, see the `tf.RaggedTensor` class documentation
+and the [Ragged Tensor Guide](/guides/ragged_tensor).
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_conversion_ops
+from tensorflow.python.ops.ragged import ragged_dispatch
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_getitem
+from tensorflow.python.ops.ragged import ragged_map_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
from tensorflow.python.ops.ragged import ragged_operators
from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.ops.ragged import segment_id_ops
-from tensorflow.python.ops.ragged.ragged_array_ops import batch_gather
-from tensorflow.python.ops.ragged.ragged_array_ops import boolean_mask
-from tensorflow.python.ops.ragged.ragged_array_ops import bounding_shape
-from tensorflow.python.ops.ragged.ragged_array_ops import concat
-from tensorflow.python.ops.ragged.ragged_array_ops import expand_dims
-from tensorflow.python.ops.ragged.ragged_array_ops import gather
-from tensorflow.python.ops.ragged.ragged_array_ops import gather_nd
-from tensorflow.python.ops.ragged.ragged_array_ops import nrows
-from tensorflow.python.ops.ragged.ragged_array_ops import row_lengths
-from tensorflow.python.ops.ragged.ragged_array_ops import row_limits
-from tensorflow.python.ops.ragged.ragged_array_ops import row_starts
-from tensorflow.python.ops.ragged.ragged_array_ops import stack
-from tensorflow.python.ops.ragged.ragged_array_ops import tile
-from tensorflow.python.ops.ragged.ragged_array_ops import value_rowids
-from tensorflow.python.ops.ragged.ragged_array_ops import where
-
-from tensorflow.python.ops.ragged.ragged_conversion_ops import from_sparse
-from tensorflow.python.ops.ragged.ragged_conversion_ops import from_tensor
-from tensorflow.python.ops.ragged.ragged_conversion_ops import to_sparse
-from tensorflow.python.ops.ragged.ragged_conversion_ops import to_tensor
-
-# pylint: disable=protected-access, wildcard-import
-from tensorflow.python.ops.ragged.ragged_elementwise_ops import *
-from tensorflow.python.ops.ragged.ragged_elementwise_ops import _symbols_to_export as _elementwise_ops
-# pylint: enable=protected-access, wildcard-import
-
-from tensorflow.python.ops.ragged.ragged_factory_ops import constant
-from tensorflow.python.ops.ragged.ragged_factory_ops import constant_value
-from tensorflow.python.ops.ragged.ragged_factory_ops import convert_to_tensor_or_ragged_tensor
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_nested_row_splits
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_nested_value_rowids
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_row_lengths
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_row_limits
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_row_splits
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_row_starts
-from tensorflow.python.ops.ragged.ragged_factory_ops import from_value_rowids
-
-from tensorflow.python.ops.ragged.ragged_functional_ops import map_inner_values
-
-from tensorflow.python.ops.ragged.ragged_map_ops import map_fn
-
-from tensorflow.python.ops.ragged.ragged_math_ops import range # pylint: disable=redefined-builtin
-
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_all
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_any
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_max
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_mean
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_min
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_prod
-from tensorflow.python.ops.ragged.ragged_math_ops import reduce_sum
-
-from tensorflow.python.ops.ragged.ragged_math_ops import segment_max
-from tensorflow.python.ops.ragged.ragged_math_ops import segment_mean
-from tensorflow.python.ops.ragged.ragged_math_ops import segment_min
-from tensorflow.python.ops.ragged.ragged_math_ops import segment_prod
-from tensorflow.python.ops.ragged.ragged_math_ops import segment_sqrt_n
-from tensorflow.python.ops.ragged.ragged_math_ops import segment_sum
-
-from tensorflow.python.ops.ragged.ragged_tensor import is_ragged
-from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
-from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensorType
-
-from tensorflow.python.ops.ragged.ragged_tensor_shape import broadcast_dynamic_shape
-from tensorflow.python.ops.ragged.ragged_tensor_shape import broadcast_to
-from tensorflow.python.ops.ragged.ragged_tensor_shape import RaggedTensorDynamicShape
-
-from tensorflow.python.ops.ragged.ragged_tensor_value import RaggedTensorValue
-
-from tensorflow.python.ops.ragged.segment_id_ops import row_splits_to_segment_ids
-from tensorflow.python.ops.ragged.segment_id_ops import segment_ids_to_row_splits
-
-from tensorflow.python.util import all_util as _all_util
-
-# Any symbol that is not referenced (with "@@name") in the module docstring
-# above, or included in the "_elementwise_ops" whitelist, will be removed.
-_all_util.remove_undocumented(__name__, _elementwise_ops)
+# Add a list of the ops that support Ragged Tensors.
+__doc__ += ragged_dispatch.ragged_op_list() # pylint: disable=redefined-builtin
diff --git a/tensorflow/python/ops/ragged/convert_to_tensor_or_ragged_tensor_op_test.py b/tensorflow/python/ops/ragged/convert_to_tensor_or_ragged_tensor_op_test.py
index ef3464f..be1ccd9 100644
--- a/tensorflow/python/ops/ragged/convert_to_tensor_or_ragged_tensor_op_test.py
+++ b/tensorflow/python/ops/ragged/convert_to_tensor_or_ragged_tensor_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.convert_to_tensor_or_ragged_tensor."""
+"""Tests for ragged_tensor.convert_to_tensor_or_ragged."""
from __future__ import absolute_import
from __future__ import division
@@ -24,12 +24,15 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedConvertToTensorOrRaggedTensorTest(test_util.TensorFlowTestCase,
- parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConvertToTensorOrRaggedTensorTest(
+ ragged_test_util.RaggedTensorTestCase, parameterized.TestCase):
#=============================================================================
# Tests where the 'value' param is a RaggedTensor
@@ -40,8 +43,8 @@
dict(pylist=[[1, 2], [3]], preferred_dtype=dtypes.string),
])
def testConvertRaggedTensor(self, pylist, dtype=None, preferred_dtype=None):
- rt = ragged.constant(pylist)
- converted = ragged.convert_to_tensor_or_ragged_tensor(
+ rt = ragged_factory_ops.constant(pylist)
+ converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
rt, dtype, preferred_dtype)
self.assertIs(converted, rt)
@@ -62,35 +65,40 @@
message,
dtype=None,
preferred_dtype=None):
- rt = ragged.constant(pylist)
+ rt = ragged_factory_ops.constant(pylist)
with self.assertRaisesRegexp(ValueError, message):
- ragged.convert_to_tensor_or_ragged_tensor(rt, dtype, preferred_dtype)
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(rt, dtype,
+ preferred_dtype)
#=============================================================================
# Tests where the 'value' param is a RaggedTensorValue
#=============================================================================
- @parameterized.parameters([
- dict(
- value=ragged.constant_value([[1, 2], [3]], dtype=np.int32),
- expected_dtype=dtypes.int32),
- dict(
- value=ragged.constant_value([[b'a', b'b'], [b'c']]),
- expected_dtype=dtypes.string),
- dict(
- value=ragged.constant_value([[1, 2], [3]], dtype=np.int32),
- dtype=dtypes.float32,
- expected_dtype=dtypes.float32),
- dict(
- value=ragged.constant_value([[1, 2], [3]], dtype=np.int32),
- preferred_dtype=dtypes.float32,
- expected_dtype=dtypes.float32),
- dict(
- value=ragged.constant_value([[1, 2], [3]], dtype=np.int32),
- preferred_dtype=dtypes.string,
- expected_dtype=dtypes.int32),
- ])
- @test_util.run_deprecated_v1
+ @parameterized.parameters(
+ [
+ dict(
+ value=ragged_factory_ops.constant_value([[1, 2], [3]],
+ dtype=np.int32),
+ expected_dtype=dtypes.int32),
+ dict(
+ value=ragged_factory_ops.constant_value([[b'a', b'b'], [b'c']]),
+ expected_dtype=dtypes.string),
+ dict(
+ value=ragged_factory_ops.constant_value([[1, 2], [3]],
+ dtype=np.int32),
+ dtype=dtypes.float32,
+ expected_dtype=dtypes.float32),
+ dict(
+ value=ragged_factory_ops.constant_value([[1, 2], [3]],
+ dtype=np.int32),
+ preferred_dtype=dtypes.float32,
+ expected_dtype=dtypes.float32),
+ dict(
+ value=ragged_factory_ops.constant_value([[1, 2], [3]],
+ dtype=np.int32),
+ preferred_dtype=dtypes.string,
+ expected_dtype=dtypes.int32),
+ ])
def testConvertRaggedTensorValue(self,
value,
dtype=None,
@@ -98,16 +106,16 @@
expected_dtype=None):
if expected_dtype is None:
expected_dtype = value.dtype if dtype is None else dtype
- converted = ragged.convert_to_tensor_or_ragged_tensor(
+ converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
value, dtype, preferred_dtype)
self.assertEqual(value.ragged_rank, converted.ragged_rank)
self.assertEqual(dtypes.as_dtype(expected_dtype), converted.dtype)
- with self.test_session():
- self.assertEqual(value.tolist(), self.evaluate(converted).tolist())
+ self.assertEqual(value.to_list(), self.eval_to_list(converted))
@parameterized.parameters([
dict(
- value=ragged.constant_value([['a', 'b'], ['c']], dtype=str),
+ value=ragged_factory_ops.constant_value([['a', 'b'], ['c']],
+ dtype=str),
dtype=dtypes.int32,
message=r"invalid literal for int\(\) with base 10: 'a'"),
])
@@ -117,7 +125,8 @@
dtype=None,
preferred_dtype=None):
with self.assertRaisesRegexp(ValueError, message):
- ragged.convert_to_tensor_or_ragged_tensor(value, dtype, preferred_dtype)
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(value, dtype,
+ preferred_dtype)
#=============================================================================
# Tests where the 'value' param is a Tensor
@@ -129,10 +138,9 @@
])
def testConvertTensor(self, pylist, dtype=None, preferred_dtype=None):
tensor = constant_op.constant(pylist)
- converted = ragged.convert_to_tensor_or_ragged_tensor(
+ converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
tensor, dtype, preferred_dtype)
- with self.test_session():
- self.assertIs(tensor, converted)
+ self.assertIs(tensor, converted)
@parameterized.parameters([
dict(
@@ -146,7 +154,6 @@
message=('Tensor conversion requested dtype string for '
'Tensor with dtype int32')),
])
- @test_util.run_deprecated_v1
def testConvertTensorError(self,
pylist,
message,
@@ -154,7 +161,8 @@
preferred_dtype=None):
tensor = constant_op.constant(pylist)
with self.assertRaisesRegexp(ValueError, message):
- ragged.convert_to_tensor_or_ragged_tensor(tensor, dtype, preferred_dtype)
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(tensor, dtype,
+ preferred_dtype)
#=============================================================================
# Tests where the 'value' param is a np.array
@@ -186,11 +194,10 @@
expected_dtype=None):
if expected_dtype is None:
expected_dtype = value.dtype if dtype is None else dtype
- converted = ragged.convert_to_tensor_or_ragged_tensor(
+ converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
value, dtype, preferred_dtype)
self.assertEqual(dtypes.as_dtype(expected_dtype), converted.dtype)
- with self.test_session():
- self.assertAllEqual(value, converted)
+ self.assertAllEqual(value, converted)
@parameterized.parameters([
dict(
@@ -204,7 +211,8 @@
dtype=None,
preferred_dtype=None):
with self.assertRaisesRegexp(ValueError, message):
- ragged.convert_to_tensor_or_ragged_tensor(value, dtype, preferred_dtype)
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(value, dtype,
+ preferred_dtype)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_array_ops.py b/tensorflow/python/ops/ragged/ragged_array_ops.py
index 603e39d..8ba8c53 100644
--- a/tensorflow/python/ops/ragged/ragged_array_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_array_ops.py
@@ -27,288 +27,18 @@
from tensorflow.python.ops import gen_ragged_array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.ragged import ragged_conversion_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_functional_ops
from tensorflow.python.ops.ragged import ragged_math_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged import ragged_util
from tensorflow.python.ops.ragged import segment_id_ops
-#===============================================================================
-# Row Partitioning
-#===============================================================================
-
-
-def value_rowids(rt_input, name=None):
- """Returns the row indices for the `values` in the given ragged tensor.
-
- `value_rowids(rt)` corresponds one-to-one with the outermost dimension of
- `rt.values`, and specifies the row containing each value. In particular,
- the row `rt[row]` consists of the values `rt.values[j]` where
- `value_rowids(rt)[j] == row`.
-
- Args:
- rt_input: The RaggedTensor whose row indices should be returned.
- name: A name prefix for the returned tensor (optional).
-
- Returns:
- A 1-D `int64` `Tensor` with shape `self.values.shape[:1]`.
- The returned tensor is nonnegative, and is sorted in ascending order.
-
- #### Example:
- ```python
- >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
- >>> rt.values.eval()
- [3, 1, 4, 1, 5, 9, 2, 6]
- >>> ragged.value_rowids(rt).eval()
- [0, 0, 0, 0, 2, 2, 2, 3] # corresponds 1:1 with rt.values
- ```
- """
- if not ragged_tensor.is_ragged(rt_input):
- raise TypeError(
- 'rt_input expected RaggedTensor, got %s' % type(rt_input).__name__)
- if (isinstance(rt_input, ragged_tensor.RaggedTensor) and
- rt_input.cached_value_rowids is not None):
- return rt_input.cached_value_rowids
-
- with ops.name_scope(name, 'RaggedValueRowIds', [rt_input]):
- return segment_id_ops.row_splits_to_segment_ids(rt_input.row_splits)
-
-
-def nrows(rt_input, out_type=dtypes.int64, name=None):
- """Returns the number of rows in the given potentially ragged tensor.
-
- I.e., the size of the outermost dimension of the tensor.
-
- Args:
- rt_input: The potentially ragged tensor whose number of rows should be
- returned.
- out_type: `dtype` for the returned tensor.
- name: A name prefix for the returned tensor (optional).
-
- Returns:
- A scalar `Tensor` with dtype `out_type`.
-
- #### Example:
- ```python
- >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
- >>> ragged.nrows(rt).eval() # rt has 5 rows.
- 5
- ```
- """
- if (isinstance(rt_input, ragged_tensor.RaggedTensor) and
- rt_input.cached_nrows is not None):
- return rt_input.cached_nrows
-
- with ops.name_scope(name, 'RaggedNRows', [rt_input]):
- if ragged_tensor.is_ragged(rt_input):
- return array_ops.shape(rt_input.row_splits, out_type=out_type)[0] - 1
- else:
- return array_ops.shape(rt_input, out_type=out_type)[0]
-
-
-def row_starts(rt_input, name=None):
- """Returns the start indices for rows in the given ragged tensor.
-
- These indices specify where the values for each row begin in
- `rt_input.values`. `ragged.row_starts(rt_input)` is equal to
- `rt_input.row_splits[:-1]`.
-
- Args:
- rt_input: The RaggedTensor whose row starts should be returned.
- name: A name prefix for the returned tensor (optional).
-
- Returns:
- A 1-D Tensor of int64 with shape `[nrows]`.
- The returned tensor is nonnegative, and is sorted in ascending order.
-
- #### Example:
- ```python
- >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
- >>> ragged.values(rt).eval()
- [3, 1, 4, 1, 5, 9, 2, 6]
- >>> ragged.row_starts(rt).eval() # indices of row starts in ragged.values
- [0, 4, 4, 7, 8]
- ```
- """
- if not ragged_tensor.is_ragged(rt_input):
- raise TypeError(
- 'rt_input expected RaggedTensor, got %s' % type(rt_input).__name__)
- with ops.name_scope(name, 'RaggedRowStarts', [rt_input]):
- return rt_input.row_splits[:-1]
-
-
-def row_limits(rt_input, name=None):
- """Returns the limit indices for rows in the given ragged tensor.
-
- These indices specify where the values for each row end in
- `rt_input.values`. `ragged.row_limits(rt_input)` is equal to
- `rt_input.row_splits[:-1]`.
-
- Args:
- rt_input: The RaggedTensor whose row limits should be returned.
- name: A name prefix for the returned tensor (optional).
-
- Returns:
- A 1-D Tensor of int64 with shape `[nrows]`.
- The returned tensor is nonnegative, and is sorted in ascending order.
-
- #### Example:
- ```python
- >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
- >>> ragged.values(rt).eval()
- [3, 1, 4, 1, 5, 9, 2, 6]
- >>> ragged.row_limits(rt).eval() # indices of row limits in ragged.values
- [4, 4, 7, 8, 8]
- ```
- """
- if not ragged_tensor.is_ragged(rt_input):
- raise TypeError(
- 'rt_input expected RaggedTensor, got %s' % type(rt_input).__name__)
- with ops.name_scope(name, 'RaggedRowLimits', [rt_input]):
- return rt_input.row_splits[1:]
-
-
-def row_lengths(rt_input, axis=1, name=None):
- """Returns the lengths of the rows in the given potentially ragged tensor.
-
- `ragged.row_lengths(rt_input)[i]` indicates the number of values in the
- `i`th row of `rt_input`.
-
- Args:
- rt_input: The potentially ragged tensor whose row lengths should be
- returned. Must have at least `axis+1` dimensions.
- axis: An integer constant indicating the axis whose row lengths should be
- returned.
- name: A name prefix for the returned tensor (optional).
-
- Returns:
- A potentially Tensor of int64 with shape `rt_input.shape[:axis]`.
-
- Raises:
- ValueError: If rt_input is a scalar, or `axis` is out of bounds.
-
- #### Example:
- ```python
- >>> rt = ragged.constant([[[3, 1, 4], [1]], [], [[5, 9], [2]], [[6]], []])
- >>> ragged.row_lengths(rt).eval() # lengths of rows in rt
- [2, 0, 2, 1, 0]
- >>> ragged.row_lengths(rt, axis=2).eval() # lengths of axis=2 rows.
- [[3, 1], [], [2, 1], [1], []]
- ```
- """
- if (isinstance(rt_input, ragged_tensor.RaggedTensor) and
- rt_input.cached_row_lengths is not None):
- return rt_input.cached_row_lengths
-
- with ops.name_scope(name, 'RaggedRowLengths', [rt_input]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
- ndims = rt_input.shape.ndims
- if ndims is not None:
- if ndims == 0:
- raise ValueError('rt_input may not be a scalar.')
- elif not -ndims <= axis < ndims:
- raise ValueError('axis=%d out of bounds: expected %d<=axis<%d.' %
- (axis, -ndims, ndims))
- if ragged_tensor.is_ragged(rt_input):
- axis = ragged_util.get_positive_axis(axis, rt_input.shape.ndims)
- if axis == 0:
- return nrows(rt_input)
- elif axis == 1:
- splits = rt_input.row_splits
- return splits[1:] - splits[:-1]
- else:
- return rt_input.with_values(row_lengths(rt_input.values, axis - 1))
- else:
- shape = array_ops.shape(rt_input, out_type=dtypes.int64)
- return array_ops.ones(shape[:axis], dtypes.int64) * shape[axis]
-
-
-def nested_row_lengths(rt_input, name=None):
- """Returns a tuple containing the row_lengths for all ragged dimensions.
-
- `nested_row_lengths(rt)` is a tuple containing the `row_lengths` tensors for
- all ragged dimensions in `rt`, ordered from outermost to innermost.
-
- Args:
- rt_input: A potentially ragged tensor.
- name: A name prefix for the returned tensors (optional).
-
- Returns:
- A `tuple` of 1-D `int64` `Tensors`. The length of the tuple is equal to
- `rt_input.ragged_rank`.
- """
- with ops.name_scope(name, 'RaggedNestedRowLengths', [rt_input]):
- rt_nested_row_lengths = []
- while isinstance(rt_input, ragged_tensor.RaggedTensor):
- rt_nested_row_lengths.append(row_lengths(rt_input))
- rt_input = rt_input.values
- return tuple(rt_nested_row_lengths)
-
-
-#===============================================================================
-# Bounding Shape
-#===============================================================================
-def bounding_shape(rt_input, axis=None, name=None):
- """Returns the tight bounding box shape for a potentially ragged tensor.
-
- Args:
- rt_input: A potentially ragged tensor.
- axis: An integer scalar or vector indicating which axes to return the
- bounding box for. If not specified, then the full bounding box is
- returned.
- name: A name prefix for the returned tensor (optional).
-
- Returns:
- An int64 `Tensor`. If `axis` is not specified, then `output`
- is a vector with `output.shape=[rt_input.shape.ndims]`. If `axis` is a
- scalar, then the `output` is a scalar. If `axis` is a vector, then
- `output` is a vector, where `output[i]` is the bounding size for
- dimension `axis[i]`.
-
- #### Example:
- ```python
- >>> rt = ragged.constant([[1, 2, 3, 4], [5], [], [6, 7, 8, 9], [10]])
- >>> ragged.bounding_shape(rt).eval().tolist()
- [5, 4]
- ```
- """
- with ops.name_scope(name, 'RaggedBoundingBox', [rt_input, axis]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
- if not ragged_tensor.is_ragged(rt_input):
- bbox = array_ops.shape(rt_input)
- return bbox if axis is None else array_ops.gather(bbox, axis)
-
- nested_splits = rt_input.nested_row_splits
- rt_inner_values = rt_input.inner_values
-
- # Optimized special cases for when axis=0 or axis=1:
- if isinstance(axis, int):
- if axis == 0:
- return array_ops.shape(nested_splits[0], out_type=dtypes.int64)[0] - 1
- elif axis == 1:
- return math_ops.maximum(math_ops.reduce_max(row_lengths(rt_input)), 0)
-
- splits_shape = array_ops.shape(rt_input.row_splits, out_type=dtypes.int64)
- inner_values_shape = array_ops.shape(rt_inner_values, out_type=dtypes.int64)
-
- ragged_dimensions = array_ops.stack([splits_shape[0] - 1] + [
- math_ops.maximum(math_ops.reduce_max(splits[1:] - splits[:-1]), 0)
- for splits in nested_splits
- ])
- inner_dimensions = inner_values_shape[1:]
-
- bbox = array_ops.concat([ragged_dimensions, inner_dimensions], axis=0)
- return bbox if axis is None else array_ops.gather(bbox, axis)
-
#===============================================================================
# ragged_gather
#===============================================================================
# TODO(edloper): Add an `axis` argument
-def gather(params, indices, name=None):
+def gather(params, indices, validate_indices=None, axis=0, name=None):
"""Gathers ragged slices from `params` axis `0` according to `indices`.
Returns `RaggedTensor` output, such that:
@@ -328,16 +58,16 @@
```python
>>> params = tf.constant(['a', 'b', 'c', 'd', 'e'])
>>> indices = tf.constant([3, 1, 2, 1, 0])
- >>> ragged_params = ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
- >>> ragged_indices = ragged.constant([[3, 1, 2], [1], [], [0]])
+ >>> ragged_params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
+ >>> ragged_indices = tf.ragged.constant([[3, 1, 2], [1], [], [0]])
- >>> print ragged.gather(params, ragged_indices).eval().tolist()
+ >>> print ragged.gather(params, ragged_indices)
[['d', 'b', 'c'], ['b'], [], ['a']]
- >>> print ragged.gather(ragged_params, indices).eval().tolist()
+ >>> print ragged.gather(ragged_params, indices)
[['e'], ['d'], [], ['d'], ['a', 'b', 'c']]
- >>> print ragged.gather(ragged_params, ragged_indices).eval().tolist()
+ >>> print ragged.gather(ragged_params, ragged_indices)
[[['e'], ['d'], []], [['d']], [], [['a', 'b', 'c']]]
```
@@ -347,6 +77,8 @@
indices: The potentially ragged tensor indicating which values to gather.
Must have dtype `int32` or `int64`. Values must be in the range `[0,
params.shape[0]]`.
+ validate_indices: Ignored.
+ axis: Must be zero.
name: A name for the operation (optional).
Returns:
@@ -357,10 +89,13 @@
Raises:
ValueError: If indices.shape.ndims is not known statically.
"""
+ del validate_indices
+ if not isinstance(axis, int) or axis != 0:
+ raise ValueError('axis>0 is not supported for ragged gather yet.')
with ops.name_scope(name, 'RaggedGather', [params, indices]):
- params = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params')
- indices = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
if ragged_tensor.is_ragged(indices):
@@ -375,13 +110,13 @@
result = gen_ragged_array_ops.ragged_gather(
indices=indices,
- params_dense_values=params.inner_values,
+ params_dense_values=params.flat_values,
params_nested_splits=params.nested_row_splits,
OUTPUT_RAGGED_RANK=indices.shape.ndims + len(params.nested_row_splits) -
1)
# Compose the RaggedTensor from splits & values.
- return ragged_factory_ops.from_nested_row_splits(
+ return ragged_tensor.RaggedTensor.from_nested_row_splits(
result.output_dense_values, result.output_nested_splits)
@@ -414,8 +149,8 @@
#### Example:
```python
- >>> params = ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
- >>> indices = ragged.constant([[1, 2, 0], [], [], [0, 0]])
+ >>> params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
+ >>> indices = tf.ragged.constant([[1, 2, 0], [], [], [0, 0]])
>>> ragged.batch_gather(params, indices)
[['b', 'c', 'a'], [], [], ['e', 'e']]
```
@@ -424,9 +159,9 @@
return array_ops.batch_gather(params, indices, name)
with ops.name_scope(name, 'RaggedBatchGather', [params, indices]):
- params = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params')
- indices = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
indices_ndims = indices.shape.ndims
if indices_ndims is None:
@@ -443,7 +178,7 @@
'not match params shape')
checks = [check_ops.assert_equal(params.row_splits, indices.row_splits)]
with ops.control_dependencies(checks):
- return ragged_factory_ops.from_row_splits(
+ return ragged_tensor.RaggedTensor.from_row_splits(
batch_gather(params.values, indices.values), indices.row_splits)
# Otherwise, indices is a 2D ragged tensor with 1 ragged dimension.
@@ -457,11 +192,11 @@
# Adjust indices from within-batch to global (in params.values), and
# then use ragged.gather to gather them.
- num_indices = row_lengths(indices)
- params_starts = row_starts(params)
+ num_indices = indices.row_lengths()
+ params_starts = params.row_starts()
adjustments = ragged_util.repeat(params_starts, num_indices, axis=0)
adjusted_index_values = math_ops.to_int64(indices.values) + adjustments
- return ragged_factory_ops.from_row_splits(
+ return ragged_tensor.RaggedTensor.from_row_splits(
gather(params.values, adjusted_index_values), indices.row_splits)
else: # params is a RaggedTensor and indices is a Tensor.
@@ -469,7 +204,7 @@
return gather(params, indices)
elif indices_ndims == 2:
# Adjust indices from batch-local to global (in params.values)
- adjustments = array_ops.expand_dims(row_starts(params), 1)
+ adjustments = array_ops.expand_dims(params.row_starts(), 1)
adjusted_indices = math_ops.to_int64(indices) + adjustments
return gather(params.values, adjusted_indices)
else:
@@ -527,9 +262,9 @@
with ops.name_scope(name, 'RaggedGatherNd', [params, indices]):
- params = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
params, name='params')
- indices = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
indices, name='indices')
indices_shape = indices.shape
indices_ndims = indices_shape.ndims
@@ -543,7 +278,7 @@
# `index_size` is the "n" in "gather_nd" -- i.e., the number of dimensions
# that each index slices into.
- index_size = indices_shape[-1].value
+ index_size = tensor_shape.dimension_value(indices_shape[-1])
if index_size is None:
raise ValueError('indices.shape[-1] must be statically known.')
@@ -555,8 +290,7 @@
if indices_is_dense:
indices = ragged_conversion_ops.from_tensor(
indices, ragged_rank=indices_ndims - 2)
- result = indices.with_inner_values(
- gather_nd(params, indices.inner_values))
+ result = indices.with_flat_values(gather_nd(params, indices.flat_values))
if (indices_is_dense and ragged_tensor.is_ragged(result) and
result.ragged_rank == indices_ndims - 2):
result = ragged_conversion_ops.to_tensor(result)
@@ -570,7 +304,7 @@
# Handle corner case: An empty index tuple selects the entire `params`
# value. So if `index_size` is zero, then tile `params`.
if index_size == 0:
- params_ndims = params.ragged_rank + array_ops.rank(params.inner_values)
+ params_ndims = params.ragged_rank + array_ops.rank(params.flat_values)
for dim in range(indices_ndims - 1):
params = expand_dims(params, axis=0)
multiples = array_ops.concat([
@@ -608,7 +342,7 @@
return array_ops.gather_nd(flattened_params, flattened_index_tuples)
flattened_index_tuples = array_ops.gather(
- row_starts(flattened_params), flattened_index_tuples)
+ flattened_params.row_starts(), flattened_index_tuples)
flattened_index_tuples += indices[..., dim]
flattened_params = flattened_params.values
@@ -704,9 +438,8 @@
"""
with ops.name_scope(name, 'RaggedMask', [data, mask]):
# Convert inputs to tensors.
- data = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- data, name='data')
- mask = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name='data')
+ mask = ragged_tensor.convert_to_tensor_or_ragged_tensor(
mask, dtypes.bool, name='mask')
# Get static rank of mask.
@@ -737,7 +470,7 @@
else:
# Count the number of True mask values in each row to find the
# lengths of the filtered rows; then convert to splits.
- int_mask = ragged_functional_ops.map_inner_values(
+ int_mask = ragged_functional_ops.map_flat_values(
math_ops.cast, mask, dtype=dtypes.int64)
masked_row_lengths = ragged_math_ops.reduce_sum(int_mask, axis=1)
splits.append(ragged_util.lengths_to_splits(masked_row_lengths))
@@ -749,7 +482,7 @@
# Add the ragged `splits` back to the result.
if keepdims:
- masked_values = ragged_factory_ops.from_nested_row_splits(
+ masked_values = ragged_tensor.RaggedTensor.from_nested_row_splits(
masked_values, splits)
return masked_values
@@ -760,7 +493,7 @@
# Get the masked splits: first get the length of each row, then filter
# out the rows that we are deleting, and convert that filtered set of
# masks back to a splits tensor.
- lengths = row_lengths(data)
+ lengths = data.row_lengths()
masked_lengths = array_ops.boolean_mask(lengths, mask)
masked_splits = ragged_util.lengths_to_splits(masked_lengths)
@@ -772,7 +505,8 @@
segment_mask = array_ops.gather(mask, segment_ids)
masked_values = boolean_mask(data.values, segment_mask, keepdims=False)
- return ragged_factory_ops.from_row_splits(masked_values, masked_splits)
+ return ragged_tensor.RaggedTensor.from_row_splits(masked_values,
+ masked_splits)
# If mask is non-ragged and has rank>1, then convert it to be ragged,
# with a ragged rank matching data.
@@ -793,7 +527,7 @@
# and values to get the innermost ragged tensor.
masked_lengths = math_ops.count_nonzero(mask, axis=-1)
flattened_masked_lengths = array_ops.reshape(masked_lengths, [-1])
- masked_values = ragged_factory_ops.from_row_lengths(
+ masked_values = ragged_tensor.RaggedTensor.from_row_lengths(
masked_values, flattened_masked_lengths)
# Wrap remaining ragged dimensions.
@@ -803,7 +537,7 @@
for dim in range(mask.shape.ndims - 3, -1, -1):
elt_size = mask_shape[dim + 1]
masked_splits = math_ops.range(split_size[dim]) * elt_size
- masked_values = ragged_factory_ops.from_row_splits(
+ masked_values = ragged_tensor.RaggedTensor.from_row_splits(
masked_values, masked_splits)
return masked_values
@@ -812,86 +546,86 @@
#===============================================================================
# Concatenation and Stacking
#===============================================================================
-def concat(rt_inputs, axis, name=None):
+def concat(values, axis, name=None):
"""Concatenates potentially ragged tensors along one dimension.
Given a list of tensors with the same rank `K` (`K >= axis`), returns a
rank-`K` `RaggedTensor` `result` such that `result[i0...iaxis]` is the
- concatenation of `[rt[i0...iaxis] for rt in rt_inputs]`.
+ concatenation of `[rt[i0...iaxis] for rt in values]`.
Args:
- rt_inputs: A list of potentially ragged tensors. May not be empty. All
- `rt_inputs` must have the same rank and the same dtype; but unlike
+ values: A list of potentially ragged tensors. May not be empty. All
+ `values` must have the same rank and the same dtype; but unlike
`tf.concat`, they can have arbitrary shapes.
axis: A python integer, indicating the dimension along which to concatenate.
(Note: Unlike `tf.concat`, the `axis` parameter must be statically known.)
Negative values are supported only if the rank of at least one
- `rt_inputs` value is statically known.
+ `values` value is statically known.
name: A name prefix for the returned tensor (optional).
Returns:
A `RaggedTensor` with rank `K`.
- `result.ragged_rank=max(axis, max(rt.ragged_rank for rt in rt_inputs]))`.
+ `result.ragged_rank=max(axis, max(rt.ragged_rank for rt in values]))`.
Raises:
- ValueError: If `rt_inputs` is empty, if `axis` is out of bounds or if
+ ValueError: If `values` is empty, if `axis` is out of bounds or if
the input tensors have different ranks.
#### Example:
```python
- >>> t1 = ragged.constant([[1, 2], [3, 4, 5]])
- >>> t2 = ragged.constant([[6], [7, 8, 9]])
+ >>> t1 = tf.ragged.constant([[1, 2], [3, 4, 5]])
+ >>> t2 = tf.ragged.constant([[6], [7, 8, 9]])
>>> ragged.concat([t1, t2], axis=0)
[[1, 2], [3, 4, 5], [6], [7, 8, 9]]
>>> ragged.concat([t1, t2], axis=1)
[[1, 2, 6], [3, 4, 5, 7, 8, 9]]
```
"""
- if not isinstance(rt_inputs, (list, tuple)):
- rt_inputs = [rt_inputs]
- with ops.name_scope(name, 'RaggedConcat', rt_inputs):
- return _ragged_stack_concat_helper(rt_inputs, axis, stack_values=False)
+ if not isinstance(values, (list, tuple)):
+ values = [values]
+ with ops.name_scope(name, 'RaggedConcat', values):
+ return _ragged_stack_concat_helper(values, axis, stack_values=False)
-def stack(rt_inputs, axis, name=None):
+def stack(values, axis=0, name=None):
"""Stacks potentially ragged tensors along one dimension.
Given a list of tensors with the same rank `K` (`K >= axis`), returns a
rank-`K+1` `RaggedTensor` `result` such that `result[i0...iaxis]` is the
- list `[rt[i0...iaxis] for rt in rt_inputs]`.
+ list `[rt[i0...iaxis] for rt in values]`.
Args:
- rt_inputs: A list of potentially ragged tensors. May not be empty. All
- `rt_inputs` must have the same rank and the same dtype; but unlike
+ values: A list of potentially ragged tensors. May not be empty. All
+ `values` must have the same rank and the same dtype; but unlike
`tf.concat`, they can have arbitrary shapes.
axis: A python integer, indicating the dimension along which to stack.
(Note: Unlike `tf.stack`, the `axis` parameter must be statically known.)
Negative values are supported only if the rank of at least one
- `rt_inputs` value is statically known.
+ `values` value is statically known.
name: A name prefix for the returned tensor (optional).
Returns:
A `RaggedTensor` with rank `K+1`.
- `result.ragged_rank=max(axis, max(rt.ragged_rank for rt in rt_inputs]))`.
+ `result.ragged_rank=max(axis, max(rt.ragged_rank for rt in values]))`.
Raises:
- ValueError: If `rt_inputs` is empty, if `axis` is out of bounds or if
+ ValueError: If `values` is empty, if `axis` is out of bounds or if
the input tensors have different ranks.
#### Example:
```python
- >>> t1 = ragged.constant([[1, 2], [3, 4, 5]])
- >>> t2 = ragged.constant([[6], [7, 8, 9]])
+ >>> t1 = tf.ragged.constant([[1, 2], [3, 4, 5]])
+ >>> t2 = tf.ragged.constant([[6], [7, 8, 9]])
>>> ragged.stack([t1, t2], axis=0)
[[[1, 2], [3, 4, 5]], [[6], [7, 9, 0]]]
>>> ragged.stack([t1, t2], axis=1)
[[[1, 2], [6]], [[3, 4, 5], [7, 8, 9]]]
```
"""
- if not isinstance(rt_inputs, (list, tuple)):
- rt_inputs = [rt_inputs]
- with ops.name_scope(name, 'RaggedConcat', rt_inputs):
- return _ragged_stack_concat_helper(rt_inputs, axis, stack_values=True)
+ if not isinstance(values, (list, tuple)):
+ values = [values]
+ with ops.name_scope(name, 'RaggedConcat', values):
+ return _ragged_stack_concat_helper(values, axis, stack_values=True)
def _ragged_stack_concat_helper(rt_inputs, axis, stack_values):
@@ -914,7 +648,7 @@
# Convert input tensors.
rt_inputs = [
- ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(
rt_input, name='rt_input') for rt_input in rt_inputs
]
@@ -965,7 +699,7 @@
values = [rt.values for rt in rt_inputs]
splits = [[rt_input.row_splits] for rt_input in rt_inputs]
with ops.control_dependencies(ragged_util.assert_splits_match(splits)):
- return ragged_factory_ops.from_row_splits(
+ return ragged_tensor.RaggedTensor.from_row_splits(
_ragged_stack_concat_helper(values, axis - 1, stack_values),
splits[0][0])
@@ -982,8 +716,8 @@
A RaggedTensor.
"""
# Concatenate the inner values together.
- inner_values = [rt.inner_values for rt in rt_inputs]
- concatenated_inner_values = array_ops.concat(inner_values, axis=0)
+ flat_values = [rt.flat_values for rt in rt_inputs]
+ concatenated_flat_values = array_ops.concat(flat_values, axis=0)
# Concatenate the splits together for each ragged dimension (adjusting
# split offsets as necessary).
@@ -997,12 +731,12 @@
# If we are performing a stack operation, then add another splits.
if stack_values:
- stack_lengths = array_ops.stack([nrows(rt) for rt in rt_inputs])
+ stack_lengths = array_ops.stack([_nrows(rt) for rt in rt_inputs])
stack_splits = ragged_util.lengths_to_splits(stack_lengths)
concatenated_nested_splits.insert(0, stack_splits)
- return ragged_factory_ops.from_nested_row_splits(concatenated_inner_values,
- concatenated_nested_splits)
+ return ragged_tensor.RaggedTensor.from_nested_row_splits(
+ concatenated_flat_values, concatenated_nested_splits)
def _ragged_stack_concat_axis_1(rt_inputs, stack_values):
@@ -1018,10 +752,10 @@
"""
num_inputs = len(rt_inputs)
- rt_nrows = nrows(rt_inputs[0])
+ rt_nrows = _nrows(rt_inputs[0])
nrows_msg = 'Input tensors have incompatible shapes.'
nrows_checks = [
- check_ops.assert_equal(nrows(rt), rt_nrows, message=nrows_msg)
+ check_ops.assert_equal(_nrows(rt), rt_nrows, message=nrows_msg)
for rt in rt_inputs[1:]
]
@@ -1045,14 +779,15 @@
# Add a new splits tensor to group together the values.
stack_splits = math_ops.range(0, rt_nrows * num_inputs + 1, num_inputs)
_copy_row_shape(rt_inputs, stack_splits)
- return ragged_factory_ops.from_row_splits(permuted_rt, stack_splits)
+ return ragged_tensor.RaggedTensor.from_row_splits(permuted_rt,
+ stack_splits)
else:
# Merge together adjacent rows by dropping the row-split indices that
# separate them.
concat_splits = permuted_rt.row_splits[::num_inputs]
_copy_row_shape(rt_inputs, concat_splits)
- return ragged_factory_ops.from_row_splits(permuted_rt.values,
- concat_splits)
+ return ragged_tensor.RaggedTensor.from_row_splits(permuted_rt.values,
+ concat_splits)
def _copy_row_shape(rt_inputs, splits):
@@ -1065,53 +800,53 @@
#===============================================================================
# Tiling
#===============================================================================
-def tile(rt_input, multiples, name=None):
+def tile(input, multiples, name=None): # pylint: disable=redefined-builtin
"""Constructs a `RaggedTensor` by tiling a given `RaggedTensor`.
- The values of `rt_input` are replicated `multiples[i]` times along the
+ The values of `input` are replicated `multiples[i]` times along the
`i`th dimension (for each dimension `i`). For every dimension `axis` in
- `rt_input`, the length of each output element in that dimension is the
+ `input`, the length of each output element in that dimension is the
length of corresponding input element multiplied by `multiples[axis]`.
Args:
- rt_input: A `RaggedTensor`.
+ input: A `RaggedTensor`.
multiples: A 1-D integer `Tensor`. Length must be the same as the number of
- dimensions in `rt_input`.
+ dimensions in `input`.
name: A name for the operation (optional).
Returns:
- A `RaggedTensor` with the same type, rank, and ragged_rank as `rt_input`.
+ A `RaggedTensor` with the same type, rank, and ragged_rank as `input`.
#### Example:
```python
- >>> rt = ragged.constant([[1, 2], [3]])
- >>> ragged.tile(rt, [3, 2]).eval().tolist()
+ >>> rt = tf.ragged.constant([[1, 2], [3]])
+ >>> ragged.tile(rt, [3, 2])
[[1, 2, 1, 2], [3, 3], [1, 2, 1, 2], [3, 3], [1, 2, 1, 2], [3, 3]]
```
"""
- with ops.name_scope(name, 'RaggedTile', [rt_input, multiples]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
+ with ops.name_scope(name, 'RaggedTile', [input, multiples]):
+ input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+ input, name='input')
multiples = ragged_util.convert_to_int_tensor(
multiples, name='multiples', dtype=dtypes.int64)
multiples.shape.assert_has_rank(1)
- if not ragged_tensor.is_ragged(rt_input):
- return array_ops.tile(rt_input, multiples, name)
+ if not ragged_tensor.is_ragged(input):
+ return array_ops.tile(input, multiples, name)
# If the constant value of `multiples` is available, then we can use it
# to skip tiling dimensions where `multiples=1`.
const_multiples = tensor_util.constant_value(multiples)
- return ragged_factory_ops.from_nested_row_splits(
- _tile_ragged_values(rt_input, multiples, const_multiples),
- _tile_ragged_splits(rt_input, multiples, const_multiples))
+ return ragged_tensor.RaggedTensor.from_nested_row_splits(
+ _tile_ragged_values(input, multiples, const_multiples),
+ _tile_ragged_splits(input, multiples, const_multiples))
def _tile_ragged_values(rt_input, multiples, const_multiples=None):
- """Builds inner_values tensor for a tiled `RaggedTensor`.
+ """Builds flat_values tensor for a tiled `RaggedTensor`.
Returns a tensor that repeats the values in
- `rt_input.inner_values` in the
+ `rt_input.flat_values` in the
appropriate pattern to construct a `RaggedTensor` that tiles `rt_input` as
specified by `multiples`.
@@ -1123,19 +858,19 @@
dimensions where `multiples=1`.
Returns:
- A `Tensor` with the same type and rank as `rt_input.inner_values`.
+ A `Tensor` with the same type and rank as `rt_input.flat_values`.
#### Example:
```python
- >>> rt = ragged.constant([[1, 2], [3]])
- >>> _tile_ragged_values(rt, [3, 2]).eval().tolist()
+ >>> rt = tf.ragged.constant([[1, 2], [3]])
+ >>> _tile_ragged_values(rt, [3, 2])
[1, 2, 1, 2, 3, 3, 1, 2, 1, 2, 3, 3, 1, 2, 1, 2, 3, 3]
```
"""
ragged_rank = rt_input.ragged_rank
nested_splits = rt_input.nested_row_splits
- # Pointers to the values in `rt_input.inner_values`.
+ # Pointers to the values in `rt_input.flat_values`.
inner_value_ids = math_ops.range(nested_splits[-1][-1])
# For each ragged dimension (working from the innermost to outermost),
@@ -1158,9 +893,9 @@
prev_splits = splits
# Gather the tiled inner values.
- ragged_tiled_values = array_ops.gather(rt_input.inner_values, inner_value_ids)
+ ragged_tiled_values = array_ops.gather(rt_input.flat_values, inner_value_ids)
- # Tile the inner_values for the uniform dimensions (i.e., for `axis=0` plus
+ # Tile the flat_values for the uniform dimensions (i.e., for `axis=0` plus
# `axis=range(ragged_rank, rank)`).
inner_repeats = array_ops.concat([multiples[:1], multiples[ragged_rank + 1:]],
axis=0)
@@ -1186,8 +921,8 @@
#### Example:
```python
- >>> rt = ragged.constant([[1, 2], [3]])
- >>> _tile_ragged_splits(rt, [3, 2]).eval().tolist()
+ >>> rt = tf.ragged.constant([[1, 2], [3]])
+ >>> _tile_ragged_splits(rt, [3, 2])
[0, 4, 6, 10, 12, 16, 18]
```
"""
@@ -1240,26 +975,26 @@
#===============================================================================
-def expand_dims(rt_input, axis, name=None):
+def expand_dims(input, axis, name=None): # pylint: disable=redefined-builtin
"""Inserts a dimension with shape 1 into a potentially ragged tensor's shape.
- Given a potentially ragged tenor `rt_input`, this operation inserts a
- dimension with size 1 at the dimension `axis` of `rt_input`'s shape.
+ Given a potentially ragged tenor `input`, this operation inserts a
+ dimension with size 1 at the dimension `axis` of `input`'s shape.
- * If `rt_input` is a `Tensor`, then this is equivalent to
+ * If `input` is a `Tensor`, then this is equivalent to
`tf.expand_dims`.
- * If `rt_input` is ragged, and `axis=0`, then the new dimension will be
+ * If `input` is ragged, and `axis=0`, then the new dimension will be
uniform; but the previously outermost dimension will become ragged.
- * If `rt_input` is ragged, and `0 < axis < rt_input.ragged_rank`, then the
+ * If `input` is ragged, and `0 < axis < input.ragged_rank`, then the
new dimension will be ragged.
- * If `rt_input` is ragged, and axis >= rt_input.ragged_rank`, then the new
+ * If `input` is ragged, and axis >= input.ragged_rank`, then the new
dimension will be uniform.
The following table gives some examples showing how `ragged.expand_dims`
impacts the shapes of different input tensors. Ragged dimensions are
indicated by enclosing them in parentheses.
- rt_input.shape | axis | result.shape
+ input.shape | axis | result.shape
----------------------- | ---- | -----------------------------
`[D1, D2]` | `0` | `[1, D1, D2]`
`[D1, D2]` | `1` | `[D1, 1, D2]`
@@ -1271,55 +1006,55 @@
`[D1, (D2), (D3), D4]` | `4` | `[D1, (D2), (D3), D4, 1]`
Args:
- rt_input: The potentially tensor that should be expanded with a new
+ input: The potentially tensor that should be expanded with a new
dimension.
axis: An integer constant indicating where the new dimension should be
inserted.
name: A name for the operation (optional).
Returns:
- A tensor with the same values as `rt_input`, with an added dimension of
+ A tensor with the same values as `input`, with an added dimension of
size 1 at `axis`.
#### Examples:
```python
- >>> rt = ragged.constant([[1, 2], [3]])
+ >>> rt = tf.ragged.constant([[1, 2], [3]])
>>> print rt.shape
TensorShape([2, None])
>>> expanded = ragged.expand_dims(rt, axis=0)
- >>> print(expanded.shape, expanded.eval().tolist())
+ >>> print(expanded.shape, expanded)
TensorShape([1, None, None]) [[[1, 2], [3]]]
>>> expanded = ragged.expand_dims(rt, axis=1)
- >>> print(expanded.shape, expanded.eval().tolist())
+ >>> print(expanded.shape, expanded)
TensorShape([2, None, None]) [[[1, 2]], [[3]]]
>>> expanded = ragged.expand_dims(rt, axis=2)
- >>> print(expanded.shape, expanded.eval().tolist())
+ >>> print(expanded.shape, expanded)
TensorShape([2, None, 1]) [[[1], [2]], [[3]]]
```
"""
- with ops.name_scope(name, 'RaggedExpandDims', [rt_input]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
+ with ops.name_scope(name, 'RaggedExpandDims', [input]):
+ input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+ input, name='input')
- if not ragged_tensor.is_ragged(rt_input):
- return array_ops.expand_dims(rt_input, axis)
+ if not ragged_tensor.is_ragged(input):
+ return array_ops.expand_dims(input, axis)
- ndims = None if rt_input.shape.ndims is None else rt_input.shape.ndims + 1
+ ndims = None if input.shape.ndims is None else input.shape.ndims + 1
axis = ragged_util.get_positive_axis(axis, ndims)
if axis == 0:
- values = rt_input
- splits = array_ops.stack([0, nrows(rt_input)])
+ values = input
+ splits = array_ops.stack([0, input.nrows()])
elif axis == 1:
- values = rt_input
- splits = math_ops.range(nrows(rt_input) + 1)
+ values = input
+ splits = math_ops.range(input.nrows() + 1)
else:
- values = expand_dims(rt_input.values, axis - 1)
- splits = rt_input.row_splits
+ values = expand_dims(input.values, axis - 1)
+ splits = input.row_splits
- return ragged_factory_ops.from_row_splits(values, splits)
+ return ragged_tensor.RaggedTensor.from_row_splits(values, splits)
#===============================================================================
@@ -1374,21 +1109,23 @@
#### Examples:
```python
>>> # Coordinates where condition is true.
- >>> condition = ragged.constant_value([[True, False, True], [False, True]])
+ >>> condition = tf.ragged.constant_value(
+ ... [[True, False, True], [False, True]])
>>> ragged.where(condition)
[[0, 0], [0, 2], [1, 1]]
>>> # Elementwise selection between x and y, based on condition.
- >>> condition = ragged.constant_value([[True, False, True], [False, True]])
- >>> x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E']])
- >>> y=ragged.constant_value([['a', 'b', 'c'], ['d', 'e']])
+ >>> condition = tf.ragged.constant_value(
+ ... [[True, False, True], [False, True]])
+ >>> x = tf.ragged.constant_value([['A', 'B', 'C'], ['D', 'E']])
+ >>> y = tf.ragged.constant_value([['a', 'b', 'c'], ['d', 'e']])
>>> ragged.where(condition, x, y)
[['A', 'b', 'C'], ['d', 'E']]
>>> # Row selection between x and y, based on condition.
>>> condition = [True, False]
- >>> x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E']])
- >>> y=ragged.constant_value([['a', 'b', 'c'], ['d', 'e']])
+ >>> x = tf.ragged.constant_value([['A', 'B', 'C'], ['D', 'E']])
+ >>> y = tf.ragged.constant_value([['a', 'b', 'c'], ['d', 'e']])
>>> ragged.where(condition, x, y)
[['A', 'B', 'C'], ['d', 'e']]
```
@@ -1396,13 +1133,13 @@
if (x is None) != (y is None):
raise ValueError('x and y must be either both None or both non-None')
with ops.name_scope('RaggedWhere', name, [condition, x, y]):
- condition = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ condition = ragged_tensor.convert_to_tensor_or_ragged_tensor(
condition, name='condition')
if x is None:
return _coordinate_where(condition)
else:
- x = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(x, name='x')
- y = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(y, name='y')
+ x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x, name='x')
+ y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y, name='y')
return _elementwise_where(condition, x, y)
@@ -1416,15 +1153,15 @@
return array_ops.where(condition, x, y)
elif condition_is_ragged and x_is_ragged and y_is_ragged:
- return ragged_functional_ops.map_inner_values(array_ops.where, condition, x,
- y)
+ return ragged_functional_ops.map_flat_values(array_ops.where, condition, x,
+ y)
elif not condition_is_ragged:
# Concatenate x and y, and then use `gather` to assemble the selected rows.
condition.shape.assert_has_rank(1)
- x_nrows = nrows(x)
+ x_nrows = _nrows(x)
x_and_y = concat([x, y], axis=0)
indices = array_ops.where(condition, math_ops.range(x_nrows),
- x_nrows + math_ops.range(nrows(y)))
+ x_nrows + math_ops.range(_nrows(y)))
return gather(x_and_y, indices)
else:
@@ -1441,7 +1178,7 @@
# Convert the first index in each coordinate to a row index and column index.
first_index = selected_coords[:, 0]
- selected_rows = array_ops.gather(value_rowids(condition), first_index)
+ selected_rows = array_ops.gather(condition.value_rowids(), first_index)
selected_row_starts = array_ops.gather(condition.row_splits, selected_rows)
selected_cols = first_index - selected_row_starts
@@ -1477,3 +1214,11 @@
pieces.append(splits[1:] + splits_offset)
splits_offset += splits[-1]
return array_ops.concat(pieces, axis=0)
+
+
+def _nrows(rt_input, out_type=dtypes.int64, name=None):
+ if isinstance(rt_input, ragged_tensor.RaggedTensor):
+ return rt_input.nrows(out_type=out_type, name=name)
+ else:
+ with ops.name_scope(name, 'RaggedNRows', [rt_input]):
+ return array_ops.shape(rt_input, out_type=out_type)[0]
diff --git a/tensorflow/python/ops/ragged/ragged_batch_gather_op_test.py b/tensorflow/python/ops/ragged/ragged_batch_gather_op_test.py
index d9d8405..431d350 100644
--- a/tensorflow/python/ops/ragged/ragged_batch_gather_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_batch_gather_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for tf.ragged.batch_gather."""
+"""Tests for ragged_array_ops.batch_gather."""
from __future__ import absolute_import
from __future__ import division
@@ -20,15 +20,20 @@
from absl.testing import parameterized
+from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedBatchGatherOpTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedBatchGatherOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
@parameterized.parameters([
@@ -37,10 +42,12 @@
#=========================================================================
dict(
descr='Docstring example',
- params=ragged.constant_value([['a', 'b', 'c'], ['d'], [], ['e']]),
- indices=ragged.constant_value([[1, 2, 0], [], [], [0, 0]]),
- expected=ragged.constant_value([[b'b', b'c', b'a'], [], [],
- [b'e', b'e']])),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d'], [],
+ ['e']]),
+ indices=ragged_factory_ops.constant_value([[1, 2, 0], [], [], [0,
+ 0]]),
+ expected=ragged_factory_ops.constant_value([[b'b', b'c', b'a'], [],
+ [], [b'e', b'e']])),
#=========================================================================
# 0 Batch Dimensions
#=========================================================================
@@ -51,9 +58,10 @@
expected=[b'd', b'c']),
dict(
descr='params: [P1, (P2)], indices: [I], result: [I, (P2)]',
- params=ragged.constant_value([['a', 'b'], [], ['c'], ['d', 'e']]),
+ params=ragged_factory_ops.constant_value([['a', 'b'], [], ['c'],
+ ['d', 'e']]),
indices=[3, 2],
- expected=ragged.constant_value([[b'd', b'e'], [b'c']])),
+ expected=ragged_factory_ops.constant_value([[b'd', b'e'], [b'c']])),
#=========================================================================
# 1 Batch Dimension
#=========================================================================
@@ -64,22 +72,24 @@
expected=[[b'c', b'a'], [b'd', b'e'], [b'h', b'g']]),
dict(
descr='params: [B1, (P1)], indices: [B1, I], result: [B1, I]',
- params=ragged.constant_value([['a', 'b', 'c'], ['d', 'e'], ['g']]),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e'],
+ ['g']]),
indices=[[2, 0], [0, 1], [0, 0]],
expected=[[b'c', b'a'], [b'd', b'e'], [b'g', b'g']]),
dict(
descr='params: [B1, P1], indices: [B1, (I)], result: [B1, (I)]',
params=[['a', 'b', 'c'], ['d', 'e', 'f'], ['g', 'h', 'i']],
- indices=ragged.constant_value([[2, 0, 2], [0], [1]]),
- expected=ragged.constant_value([[b'c', b'a', b'c'], [b'd'], [b'h']])),
+ indices=ragged_factory_ops.constant_value([[2, 0, 2], [0], [1]]),
+ expected=ragged_factory_ops.constant_value([[b'c', b'a', b'c'],
+ [b'd'], [b'h']])),
dict(
descr=('params: [B1, (P1), (P2), P3], indices: [B1, I], '
'result: [B1, I, (P2), P3]'),
- params=ragged.constant_value(
+ params=ragged_factory_ops.constant_value(
[[[['a']], [['b'], ['c']]], [[['d'], ['e']], [['f']]], [[['g']]]],
ragged_rank=2),
indices=[[1, 0], [0, 1], [0, 0]],
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[[b'b'], [b'c']], [[b'a']]], [[[b'd'], [b'e']], [[b'f']]],
[[[b'g']], [[b'g']]]],
ragged_rank=2)),
@@ -95,31 +105,31 @@
dict(
descr=('params: [B1, (B2), P1], indices: [B1, (B2), I], '
'result: [B1, (B2), I]'),
- params=ragged.constant_value(
+ params=ragged_factory_ops.constant_value(
[[['a', 'b', 'c'], ['d', 'e', 'f']], [['g', 'h', 'i']]],
ragged_rank=1),
- indices=ragged.constant_value([[[2, 0], [0, 1]], [[1, 0]]],
- ragged_rank=1),
- expected=ragged.constant_value(
+ indices=ragged_factory_ops.constant_value(
+ [[[2, 0], [0, 1]], [[1, 0]]], ragged_rank=1),
+ expected=ragged_factory_ops.constant_value(
[[[b'c', b'a'], [b'd', b'e']], [[b'h', b'g']]], ragged_rank=1)),
dict(
descr=('params: [B1, (B2), (P1)], indices: [B1, (B2), I], '
'result: [B1, (B2), I]'),
- params=ragged.constant_value([[['a', 'b', 'c'], ['d']], [['e', 'f']]],
- ragged_rank=2),
- indices=ragged.constant_value([[[2, 0], [0, 0]], [[1, 0]]],
- ragged_rank=1),
- expected=ragged.constant_value(
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d']], [['e', 'f']]], ragged_rank=2),
+ indices=ragged_factory_ops.constant_value(
+ [[[2, 0], [0, 0]], [[1, 0]]], ragged_rank=1),
+ expected=ragged_factory_ops.constant_value(
[[[b'c', b'a'], [b'd', b'd']], [[b'f', b'e']]], ragged_rank=1)),
dict(
descr=('params: [B1, (B2), P1], indices: [B1, (B2), (I)], '
'result: [B1, (B2), (I)]'),
- params=ragged.constant_value(
+ params=ragged_factory_ops.constant_value(
[[['a', 'b', 'c'], ['d', 'e', 'f']], [['g', 'h', 'i']]],
ragged_rank=1),
- indices=ragged.constant_value([[[2, 1, 0], [0]], [[1, 1]]],
- ragged_rank=2),
- expected=ragged.constant_value(
+ indices=ragged_factory_ops.constant_value(
+ [[[2, 1, 0], [0]], [[1, 1]]], ragged_rank=2),
+ expected=ragged_factory_ops.constant_value(
[[[b'c', b'b', b'a'], [b'd']], [[b'h', b'h']]], ragged_rank=2)),
#=========================================================================
# 3 Batch Dimensions
@@ -128,74 +138,77 @@
descr=(
'params: [B1, (B2), (B3), (P1)], indices: [B1, (B2), (B3), I], '
'result: [B1, (B2), (B3), I]'),
- params=ragged.constant_value(
+ params=ragged_factory_ops.constant_value(
[[[['a', 'b', 'c'], ['d']], [['e', 'f']]]], ragged_rank=3),
- indices=ragged.constant_value([[[[2, 0], [0, 0]], [[1, 0]]]],
- ragged_rank=2),
- expected=ragged.constant_value(
+ indices=ragged_factory_ops.constant_value(
+ [[[[2, 0], [0, 0]], [[1, 0]]]], ragged_rank=2),
+ expected=ragged_factory_ops.constant_value(
[[[[b'c', b'a'], [b'd', b'd']], [[b'f', b'e']]]], ragged_rank=2)),
])
- @test_util.run_deprecated_v1
def testRaggedBatchGather(self, descr, params, indices, expected):
- result = ragged.batch_gather(params, indices)
- self.assertEqual(
- getattr(result, 'ragged_rank', 0), getattr(expected, 'ragged_rank', 0))
- with self.test_session():
- if hasattr(expected, 'tolist'):
- expected = expected.tolist()
- self.assertEqual(result.eval().tolist(), expected)
+ result = ragged_array_ops.batch_gather(params, indices)
+ self.assertRaggedEqual(result, expected)
- @test_util.run_deprecated_v1
def testRaggedBatchGatherUnknownRankError(self):
+ if context.executing_eagerly():
+ return
params = [['a', 'b'], ['c', 'd']]
indices = array_ops.placeholder(dtypes.int32, shape=None)
- ragged_indices = ragged.from_row_splits(indices, [0, 2, 4])
+ ragged_indices = ragged_tensor.RaggedTensor.from_row_splits(
+ indices, [0, 2, 4])
with self.assertRaisesRegexp(
ValueError, 'batch_gather does not allow indices with unknown shape.'):
- ragged.batch_gather(params, indices)
+ ragged_array_ops.batch_gather(params, indices)
with self.assertRaisesRegexp(
ValueError, 'batch_gather does not allow indices with unknown shape.'):
- ragged.batch_gather(params, ragged_indices)
+ ragged_array_ops.batch_gather(params, ragged_indices)
- @parameterized.parameters([
- dict(
- params=ragged.constant([['a'], ['b'], ['c']]),
- indices=ragged.constant([[0], [0]]),
- message='Dimensions 3 and 2 are not compatible'),
- dict(
- params=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
- indices=ragged.constant([[[0, 0], [0, 0, 0]], [[0]]]),
- message='batch shape from indices does not match params shape'),
- dict(
- params=ragged.constant([[[0, 0], [0, 0, 0]], [[0]]]),
- indices=ragged.constant([[[0, 0]], [[0, 0, 0]], [[0]]]),
- message='Dimensions must be equal, but are 3 and 4'),
- dict(
- params=ragged.constant([[[0, 0], [0, 0, 0]], [[0]], [[0]]]),
- indices=ragged.constant([[[0, 0]], [[0, 0, 0]], [[0]]]),
- error=errors.InvalidArgumentError,
- message='Condition x == y did not hold element-wise'),
- dict(
- params=ragged.constant(['a', 'b', 'c']),
- indices=ragged.constant([[0], [0]]),
- message='batch shape from indices does not match params shape'),
- dict(params=ragged.constant_value([['a']]),
- indices=0,
- message='indices.rank must be at least 1.'),
- dict(params=ragged.constant_value([['a']]),
- indices=[[[0]]],
- message='batch shape from indices does not match params shape'),
- ])
- @test_util.run_deprecated_v1
+ @parameterized.parameters(
+ [
+ dict(
+ params=ragged_factory_ops.constant_value([['a'], ['b'], ['c']]),
+ indices=ragged_factory_ops.constant_value([[0], [0]]),
+ message='Dimensions 3 and 2 are not compatible'),
+ dict(
+ params=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+ indices=ragged_factory_ops.constant_value([[[0, 0], [0, 0, 0]],
+ [[0]]]),
+ message='batch shape from indices does not match params shape'),
+ dict( # rank mismatch
+ params=ragged_factory_ops.constant_value([[[0, 0], [0, 0, 0]],
+ [[0]]]),
+ indices=ragged_factory_ops.constant_value([[[0, 0]], [[0, 0, 0]],
+ [[0]]]),
+ error=(ValueError, errors.InvalidArgumentError)),
+ dict(
+ params=ragged_factory_ops.constant_value([[[0, 0], [0, 0, 0]],
+ [[0]], [[0]]]),
+ indices=ragged_factory_ops.constant_value([[[0, 0]], [[0, 0, 0]],
+ [[0]]]),
+ error=errors.InvalidArgumentError,
+ message='.*Condition x == y did not hold.*'),
+ dict(
+ params=ragged_factory_ops.constant_value(['a', 'b', 'c']),
+ indices=ragged_factory_ops.constant_value([[0], [0]]),
+ message='batch shape from indices does not match params shape'),
+ dict(
+ params=ragged_factory_ops.constant_value([['a']]),
+ indices=0,
+ message='indices.rank must be at least 1.'),
+ dict(
+ params=ragged_factory_ops.constant_value([['a']]),
+ indices=[[[0]]],
+ message='batch shape from indices does not match params shape'),
+ ])
def testRaggedBatchGatherStaticError(self,
params,
indices,
- message,
+ message=None,
error=ValueError):
with self.assertRaisesRegexp(error, message):
- ragged.batch_gather(params, indices)
+ ragged_array_ops.batch_gather(params, indices)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_boolean_mask_op_test.py b/tensorflow/python/ops/ragged/ragged_boolean_mask_op_test.py
index d939d9d..19f7d21 100644
--- a/tensorflow/python/ops/ragged/ragged_boolean_mask_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_boolean_mask_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.boolean_mask."""
+"""Tests for ragged_array_ops.boolean_mask."""
from __future__ import absolute_import
from __future__ import division
@@ -20,15 +20,19 @@
from absl.testing import parameterized
+from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedBooleanMaskOpTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedBooleanMaskOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
# Define short constants for true & false, so the data & mask can be lined
# up in the examples below. This makes it easier to read the examples, to
@@ -51,25 +55,25 @@
data=[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
mask=[[T, F, T], [F, F, F], [T, F, F]],
keepdims=True,
- expected=ragged.constant_value([[1, 3], [], [7]])),
+ expected=ragged_factory_ops.constant_value([[1, 3], [], [7]])),
dict(
descr='Docstring example 3',
- data=ragged.constant_value([[1, 2, 3], [4], [5, 6]]),
- mask=ragged.constant_value([[F, F, T], [F], [T, T]]),
+ data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [5, 6]]),
+ mask=ragged_factory_ops.constant_value([[F, F, T], [F], [T, T]]),
keepdims=False,
expected=[3, 5, 6]),
dict(
descr='Docstring example 4',
- data=ragged.constant_value([[1, 2, 3], [4], [5, 6]]),
- mask=ragged.constant_value([[F, F, T], [F], [T, T]]),
+ data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [5, 6]]),
+ mask=ragged_factory_ops.constant_value([[F, F, T], [F], [T, T]]),
keepdims=True,
- expected=ragged.constant_value([[3], [], [5, 6]])),
+ expected=ragged_factory_ops.constant_value([[3], [], [5, 6]])),
dict(
descr='Docstring example 5',
- data=ragged.constant_value([[1, 2, 3], [4], [5, 6]]),
+ data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [5, 6]]),
mask=[True, False, True],
keepdims=False,
- expected=ragged.constant_value([[1, 2, 3], [5, 6]])),
+ expected=ragged_factory_ops.constant_value([[1, 2, 3], [5, 6]])),
#=========================================================================
# Uniform data and uniform mask.
#=========================================================================
@@ -90,7 +94,8 @@
data=[[1, 2, 3], [4, 5, 6], [7, 8, 9], [0, 1, 2], [3, 4, 5]],
mask=[[F, F, F], [T, F, T], [T, T, T], [F, F, F], [T, T, F]],
keepdims=True,
- expected=ragged.constant_value([[], [4, 6], [7, 8, 9], [], [3, 4]])),
+ expected=ragged_factory_ops.constant_value(
+ [[], [4, 6], [7, 8, 9], [], [3, 4]])),
dict(
descr='data.shape=[3, 2, 2]; mask.shape=[3]; keepdims=True',
data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
@@ -108,8 +113,9 @@
data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
mask=[[T, F], [T, T], [F, F]],
keepdims=True,
- expected=ragged.constant_value([[[1, 2]], [[5, 6], [7, 8]], []],
- ragged_rank=1)),
+ expected=ragged_factory_ops.constant_value(
+ [[[1, 2]], [[5, 6], [7, 8]], []],
+ ragged_rank=1)),
dict(
descr='data.shape=[3, 2, 2]; mask.shape=[3, 2]; keepdims=False',
data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
@@ -121,7 +127,7 @@
data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
mask=[[[T, T], [F, T]], [[F, F], [F, F]], [[T, F], [T, T]]],
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2], [4]], [[], []], [[2], [6, 8]]])),
dict(
descr='data.shape=mask.shape=[2, 2, 2, 2]; keepdims=True',
@@ -130,7 +136,7 @@
mask=[[[[T, T], [F, F]], [[T, F], [F, F]]],
[[[F, F], [F, F]], [[T, T], [T, F]]]],
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[[1, 2], []], [[5], []]], [[[], []], [[1, 3], [5]]]])),
dict(
descr='data.shape=mask.shape=[2, 2, 2, 2]; keepdims=False',
@@ -146,63 +152,64 @@
#=========================================================================
dict(
descr='data.shape=[5, (D2)]; mask.shape=[5, (D2)]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[1, 2], [3, 4, 5, 6], [7, 8, 9], [], [1, 2, 3]]),
- mask=ragged.constant_value(
+ mask=ragged_factory_ops.constant_value(
[[F, F], [F, T, F, T], [F, F, F], [], [T, F, T]]),
keepdims=True,
- expected=ragged.constant_value([[], [4, 6], [], [], [1, 3]])),
+ expected=ragged_factory_ops.constant_value(
+ [[], [4, 6], [], [], [1, 3]])),
dict(
descr='data.shape=[3, (D2), (D3)]; mask.shape=[3, (D2)]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]]),
- mask=ragged.constant_value([[T, F], [T, T], [F, F]]),
+ mask=ragged_factory_ops.constant_value([[T, F], [T, T], [F, F]]),
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2]], [[5, 6], [7, 8]], []])),
dict(
descr='data.shape=[3, (D2), (D3)]; mask.shape=[3, (D2)]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]]),
- mask=ragged.constant_value([[T, F], [T, T], [F, F]]),
+ mask=ragged_factory_ops.constant_value([[T, F], [T, T], [F, F]]),
keepdims=False,
- expected=ragged.constant_value([[1, 2], [5, 6], [7, 8]])),
+ expected=ragged_factory_ops.constant_value([[1, 2], [5, 6], [7, 8]])),
dict(
descr='data.shape=[3, (D2), D3]; mask.shape=[3, (D2)]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[5, 6], [7, 8], [2, 4]], [[6, 8]]],
ragged_rank=1),
- mask=ragged.constant_value([[T, F], [T, T, F], [F]]),
+ mask=ragged_factory_ops.constant_value([[T, F], [T, T, F], [F]]),
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2]], [[5, 6], [7, 8]], []],
ragged_rank=1)),
dict(
descr='data.shape=[3, (D2), D3]; mask.shape=[3, (D2)]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
ragged_rank=1),
- mask=ragged.constant_value([[T, F], [T, T], [F, F]]),
+ mask=ragged_factory_ops.constant_value([[T, F], [T, T], [F, F]]),
keepdims=False,
expected=[[1, 2], [5, 6], [7, 8]]),
dict(
descr='data.shape=[3, (D2), (D3)]; mask.shape=[3, (D2), (D3)]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4]]]),
- mask=ragged.constant_value(
+ mask=ragged_factory_ops.constant_value(
[[[T, T], [F, T]], [[F, F], [F, F]], [[T, F]]]),
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2], [4]], [[], []], [[2]]])),
dict(
descr=('data.shape=[3, (D2), (D3), (D4)]; '
'mask.shape=[3, (D2), (D3), (D4)]'),
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[[1, 2], [3, 4]], [[5, 6]]], [[[2, 4], [6, 8]]]]),
- mask=ragged.constant_value(
+ mask=ragged_factory_ops.constant_value(
[[[[T, T], [F, F]], [[T, F]]], [[[F, F], [T, T]]]]),
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[[1, 2], []], [[5]]], [[[], [6, 8]]]])),
#=========================================================================
@@ -211,142 +218,132 @@
dict(
descr='data.shape=[2, 3]; mask.shape=[2, (3)]',
data=[[1, 2, 3], [4, 5, 6]],
- mask=ragged.constant_value([[T, F, F], [F, T, T]]),
+ mask=ragged_factory_ops.constant_value([[T, F, F], [F, T, T]]),
keepdims=True,
- expected=ragged.constant_value([[1], [5, 6]])),
+ expected=ragged_factory_ops.constant_value([[1], [5, 6]])),
dict(
descr='data.shape=[2, 3, 2]; mask.shape=[2, (3)]',
data=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [2, 4]]],
- mask=ragged.constant_value([[T, F, F], [F, T, T]]),
+ mask=ragged_factory_ops.constant_value([[T, F, F], [F, T, T]]),
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2]], [[9, 0], [2, 4]]],
ragged_rank=1)),
dict(
descr='data.shape=[2, 3, 2]; mask.shape=[2, (3), 2]',
data=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [2, 4]]],
- mask=ragged.constant_value(
+ mask=ragged_factory_ops.constant_value(
[[[T, F], [F, F], [T, T]], [[T, F], [F, T], [F, F]]],
ragged_rank=1),
keepdims=True,
- expected=ragged.constant_value([[[1], [], [5, 6]], [[7], [0], []]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[1], [], [5, 6]], [[7], [0], []]])),
#=========================================================================
# Ragged data and uniform mask.
#=========================================================================
dict(
descr='data.shape=[4, (D2)]; mask.shape=[4]',
- data=ragged.constant_value([[1, 2, 3], [4], [], [5, 6]]),
+ data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [], [5, 6]]),
mask=[T, F, T, F],
keepdims=False,
- expected=ragged.constant_value([[1, 2, 3], []])),
+ expected=ragged_factory_ops.constant_value([[1, 2, 3], []])),
dict(
descr='data.shape=[4, (D2), (D3)]; mask.shape=[4]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2, 3]], [[4], []], [[5, 6]], []]),
mask=[T, F, T, T],
keepdims=False,
- expected=ragged.constant_value([[[1, 2, 3]], [[5, 6]], []])),
+ expected=ragged_factory_ops.constant_value(
+ [[[1, 2, 3]], [[5, 6]], []])),
dict(
descr='data.shape=[4, (D2), 2]; mask.shape=[4]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [], [[5, 6]], [[7, 8], [9, 0], [1, 2]]],
ragged_rank=1),
mask=[T, F, F, T],
keepdims=False,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[7, 8], [9, 0], [1, 2]]],
ragged_rank=1)),
dict(
descr='data.shape=[4, (D2), 2]; mask.shape=[4]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [], [[5, 6]], [[7, 8], [9, 0], [1, 2]]],
ragged_rank=1),
mask=[T, F, F, T],
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2], [3, 4]], [[7, 8], [9, 0], [1, 2]]],
ragged_rank=1)),
dict(
descr='data.shape=[1, (2)]; mask.shape=[1, 2]',
- data=ragged.constant_value([[1, 2]]),
+ data=ragged_factory_ops.constant_value([[1, 2]]),
mask=[[T, F]],
keepdims=True,
- expected=ragged.constant_value([[1]])),
+ expected=ragged_factory_ops.constant_value([[1]])),
dict(
descr='data.shape=[2, (2), (D3)]; mask.shape=[2, 2]',
- data=ragged.constant_value([[[1], [2, 3]], [[], [4, 5, 6]]]),
+ data=ragged_factory_ops.constant_value(
+ [[[1], [2, 3]], [[], [4, 5, 6]]]),
mask=[[T, F], [T, T]],
keepdims=True,
- expected=ragged.constant_value([[[1]], [[], [4, 5, 6]]])),
+ expected=ragged_factory_ops.constant_value([[[1]], [[], [4, 5, 6]]])),
dict(
descr='data.shape=[2, (2), 3]; mask.shape=[2, 2]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]],
ragged_rank=1),
mask=[[T, F], [T, T]],
keepdims=True,
- expected=ragged.constant_value(
+ expected=ragged_factory_ops.constant_value(
[[[1, 2, 3]], [[7, 8, 9], [2, 4, 6]]],
ragged_rank=1)),
dict(
descr='data.shape=[2, (2), 3]; mask.shape=[2, 2, 3]',
- data=ragged.constant_value(
+ data=ragged_factory_ops.constant_value(
[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]],
ragged_rank=1),
mask=[[[T, F, F], [T, F, T]], [[T, F, T], [F, F, F]]],
keepdims=True,
- expected=ragged.constant_value([[[1], [4, 6]], [[7, 9], []]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[1], [4, 6]], [[7, 9], []]])),
]) # pyformat: disable
- @test_util.run_deprecated_v1
def testBooleanMask(self, descr, data, mask, keepdims, expected):
- actual = ragged.boolean_mask(data, mask, keepdims=keepdims)
- self.assertEqual(
- getattr(actual, 'ragged_rank', 0), getattr(expected, 'ragged_rank', 0))
- with self.test_session():
- if isinstance(expected, ragged.RaggedTensorValue):
- expected = expected.tolist()
- self.assertEqual(actual.eval().tolist(), expected)
+ actual = ragged_array_ops.boolean_mask(data, mask, keepdims=keepdims)
+ self.assertRaggedEqual(actual, expected)
- @test_util.run_deprecated_v1
def testErrors(self):
- self.assertRaisesRegexp(ValueError,
- r'mask\.shape\.ndims must be kown statically',
- ragged.boolean_mask, [[1, 2]],
- array_ops.placeholder(dtypes.bool))
+ if not context.executing_eagerly():
+ self.assertRaisesRegexp(ValueError,
+ r'mask\.shape\.ndims must be kown statically',
+ ragged_array_ops.boolean_mask, [[1, 2]],
+ array_ops.placeholder(dtypes.bool))
- self.assertRaisesRegexp(TypeError,
- "Expected bool, got 0 of type 'int' instead.",
- ragged.boolean_mask, [[1, 2]], [[0, 1]])
+ self.assertRaises(TypeError, ragged_array_ops.boolean_mask, [[1, 2]],
+ [[0, 1]])
self.assertRaisesRegexp(
ValueError, 'Tensor conversion requested dtype bool for '
- 'RaggedTensor with dtype int32', ragged.boolean_mask,
- ragged.constant([[1, 2]]), ragged.constant([[0, 0]]))
+ 'RaggedTensor with dtype int32', ragged_array_ops.boolean_mask,
+ ragged_factory_ops.constant([[1, 2]]),
+ ragged_factory_ops.constant([[0, 0]]))
self.assertRaisesRegexp(
ValueError, r'Shapes \(1, 2\) and \(1, 3\) are incompatible',
- ragged.boolean_mask, [[1, 2]], [[True, False, True]])
-
- # self.assertRaisesRegexp(ValueError,
- # r'data=.* is non-ragged but mask=.* is ragged',
- # ragged.boolean_mask, [[1, 2]],
- # ragged.constant([[True, False]]))
-
- # self.assertRaisesRegexp(
- # ValueError, r'data=.* is ragged but mask=.* is non-ragged',
- # ragged.boolean_mask, ragged.constant([[1, 2]]), [[True, False]])
+ ragged_array_ops.boolean_mask, [[1, 2]], [[True, False, True]])
self.assertRaisesRegexp(errors.InvalidArgumentError,
r'Inputs must have identical ragged splits',
- ragged.boolean_mask, ragged.constant([[1, 2]]),
- ragged.constant([[True, False, True]]))
+ ragged_array_ops.boolean_mask,
+ ragged_factory_ops.constant([[1, 2]]),
+ ragged_factory_ops.constant([[True, False, True]]))
self.assertRaisesRegexp(ValueError, 'mask cannot be scalar',
- ragged.boolean_mask, [[1, 2]], True)
+ ragged_array_ops.boolean_mask, [[1, 2]], True)
- self.assertRaisesRegexp(ValueError,
- 'mask cannot be scalar', ragged.boolean_mask,
- ragged.constant([[1, 2]]), True)
+ self.assertRaisesRegexp(ValueError, 'mask cannot be scalar',
+ ragged_array_ops.boolean_mask,
+ ragged_factory_ops.constant([[1, 2]]), True)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_concat_op_test.py b/tensorflow/python/ops/ragged/ragged_concat_op_test.py
index 3699f90..254afda 100644
--- a/tensorflow/python/ops/ragged/ragged_concat_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_concat_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.concat."""
+"""Tests for ragged_array_ops.concat."""
from __future__ import absolute_import
from __future__ import division
@@ -20,22 +20,27 @@
from absl.testing import parameterized
+from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedConcatOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConcatOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
def _rt_inputs_to_tensors(self, rt_inputs, ragged_ranks=None):
if ragged_ranks is None:
ragged_ranks = [None] * len(rt_inputs)
- return [
- ragged.constant(rt_input, ragged_rank=rrank)
+ return [ # pylint: disable=g-long-ternary
+ ragged_factory_ops.constant(rt_input, ragged_rank=rrank)
if rrank != 0 else constant_op.constant(rt_input)
for (rt_input, rrank) in zip(rt_inputs, ragged_ranks)
]
@@ -221,7 +226,6 @@
axis=0,
expected=[[b'a00', b'a01'], [], [b'a20', b'a21']]),
) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedConcat(self,
descr,
rt_inputs,
@@ -231,13 +235,12 @@
expected_ragged_rank=None,
expected_shape=None):
rt_inputs = self._rt_inputs_to_tensors(rt_inputs, ragged_ranks)
- concatenated = ragged.concat(rt_inputs, axis)
+ concatenated = ragged_array_ops.concat(rt_inputs, axis)
if expected_ragged_rank is not None:
self.assertEqual(concatenated.ragged_rank, expected_ragged_rank)
if expected_shape is not None:
self.assertEqual(concatenated.shape.as_list(), expected_shape)
- with self.test_session():
- self.assertEqual(concatenated.eval().tolist(), expected)
+ self.assertRaggedEqual(concatenated, expected)
@parameterized.parameters(
dict(
@@ -264,13 +267,17 @@
ragged_ranks=(0, 0),
rt_inputs=([[1, 2]], [[3, 4], [5, 6]]),
axis=1,
- error=ValueError,
- message='Dimension 0 in both shapes must be equal'),
+ error=(ValueError, errors.InvalidArgumentError)),
)
- @test_util.run_deprecated_v1
- def testStaticError(self, rt_inputs, axis, error, message, ragged_ranks=None):
+ def testStaticError(self,
+ rt_inputs,
+ axis,
+ error,
+ message=None,
+ ragged_ranks=None):
rt_inputs = self._rt_inputs_to_tensors(rt_inputs, ragged_ranks)
- self.assertRaisesRegexp(error, message, ragged.concat, rt_inputs, axis)
+ self.assertRaisesRegexp(error, message, ragged_array_ops.concat, rt_inputs,
+ axis)
@parameterized.parameters([
dict(
@@ -280,27 +287,28 @@
error=errors.InvalidArgumentError,
message='Input tensors have incompatible shapes'),
])
- @test_util.run_deprecated_v1
def testRuntimeError(self, rt_inputs, axis, error, message,
ragged_ranks=None):
+ if context.executing_eagerly():
+ return
rt_inputs = [
array_ops.placeholder_with_default(rt, shape=None) for rt in rt_inputs
]
- concatenated = ragged.concat(rt_inputs, axis)
- with self.test_session():
- self.assertRaisesRegexp(error, message, concatenated.eval)
+ concatenated = ragged_array_ops.concat(rt_inputs, axis)
+ with self.assertRaisesRegexp(error, message):
+ self.evaluate(concatenated)
- @test_util.run_deprecated_v1
def testNegativeAxisWithUnknownRankError(self):
+ if context.executing_eagerly():
+ return
rt_inputs = [
array_ops.placeholder(dtypes.int64),
array_ops.placeholder(dtypes.int64)
]
self.assertRaisesRegexp(
ValueError, r'axis may only be negative if ndims is statically known.',
- ragged.concat, rt_inputs, -1)
+ ragged_array_ops.concat, rt_inputs, -1)
- @test_util.run_deprecated_v1
def testSingleTensorInput(self):
"""Tests ragged_concat with a single tensor input.
@@ -308,10 +316,9 @@
also pass in a single value (as with tf.concat), in which case it simply
returns that tensor. This test exercises that path.
"""
- rt_inputs = ragged.constant([[1, 2], [3, 4]])
- concatenated = ragged.concat(rt_inputs, 0)
- with self.test_session():
- self.assertEqual(concatenated.eval().tolist(), [[1, 2], [3, 4]])
+ rt_inputs = ragged_factory_ops.constant([[1, 2], [3, 4]])
+ concatenated = ragged_array_ops.concat(rt_inputs, 0)
+ self.assertRaggedEqual(concatenated, [[1, 2], [3, 4]])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_const_op_test.py b/tensorflow/python/ops/ragged/ragged_const_op_test.py
index 2505b23..29a9bdf 100644
--- a/tensorflow/python/ops/ragged/ragged_const_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_const_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.constant."""
+"""Tests for ragged_factory_ops.constant."""
from __future__ import absolute_import
from __future__ import division
@@ -20,15 +20,18 @@
from absl.testing import parameterized
-
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
+from tensorflow.python.ops import ragged
from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedConstOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConstOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters(
#=========================================================================
@@ -133,7 +136,6 @@
dict(pylist=[[b'a', b'b'], [b'c'], [b'd', b'e', b'f']],
dtype=dtypes.string),
)
- @test_util.run_deprecated_v1
def testRaggedConst(self,
pylist,
dtype=None,
@@ -176,23 +178,14 @@
# If inner_shape was explicitly specified, check it.
if inner_shape is not None:
if isinstance(rt, ragged_tensor.RaggedTensor):
- self.assertEqual(rt.inner_values.shape.as_list()[1:], list(inner_shape))
+ self.assertEqual(rt.flat_values.shape.as_list()[1:], list(inner_shape))
else:
self.assertEqual(rt.shape.as_list(), list(inner_shape))
if expected_shape is not None:
self.assertEqual(tuple(rt.shape.as_list()), expected_shape)
- with self.test_session():
- result = self.evaluate(rt)
- if rt.shape.ndims > 0:
- self.assertEqual(result.tolist(), pylist)
- if expected_shape is not None:
- self.assertEqual(result.shape, expected_shape)
- else:
- self.assertEqual(result, pylist)
- if expected_shape is not None:
- self.assertEqual((), expected_shape)
+ self.assertRaggedEqual(rt, pylist)
@parameterized.parameters(
dict(
@@ -236,11 +229,7 @@
exception=ValueError,
message='pylist has scalar values depth 2, but ragged_rank=2 '
'requires scalar value depth greater than 2'),
- dict(
- pylist=[1, 2, 3],
- inner_shape=(1, 1),
- exception=TypeError,
- message='Expected Tensor\'s shape'),
+ dict(pylist=[1, 2, 3], inner_shape=(1, 1), exception=TypeError),
dict(
pylist=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
inner_shape=(2, 2),
@@ -259,7 +248,6 @@
exception=ValueError,
message='inner values have inconsistent shape'),
)
- @test_util.run_deprecated_v1
def testRaggedConstError(self,
pylist,
dtype=None,
@@ -308,9 +296,9 @@
message=None):
"""Tests for the _find_scalar_and_max_depth helper function."""
if exception is not None:
- self.assertRaisesRegexp(
- exception, message,
- ragged_factory_ops._find_scalar_and_max_depth, pylist)
+ self.assertRaisesRegexp(exception, message,
+ ragged_factory_ops._find_scalar_and_max_depth,
+ pylist)
else:
self.assertEqual(
ragged_factory_ops._find_scalar_and_max_depth(pylist),
@@ -360,11 +348,11 @@
if exception is not None:
self.assertRaisesRegexp(
exception, message,
- ragged_factory_ops._default_inner_shape_for_pylist, pylist,
+ ragged.ragged_factory_ops._default_inner_shape_for_pylist, pylist,
ragged_rank)
else:
self.assertEqual(
- ragged_factory_ops._default_inner_shape_for_pylist(
+ ragged.ragged_factory_ops._default_inner_shape_for_pylist(
pylist, ragged_rank), inner_shape)
diff --git a/tensorflow/python/ops/ragged/ragged_constant_value_op_test.py b/tensorflow/python/ops/ragged/ragged_constant_value_op_test.py
index d805189..7f47459 100644
--- a/tensorflow/python/ops/ragged/ragged_constant_value_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_constant_value_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.constant_value."""
+"""Tests for ragged_factory_ops.constant_value."""
from __future__ import absolute_import
from __future__ import division
@@ -22,11 +22,14 @@
import numpy as np
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedConstantValueOpTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConstantValueOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
@parameterized.parameters(
@@ -144,8 +147,8 @@
inner_shape=None,
expected_shape=None,
expected_dtype=None):
- """Tests that `ragged_value(pylist).tolist() == pylist`."""
- rt = ragged.constant_value(
+ """Tests that `ragged_value(pylist).to_list() == pylist`."""
+ rt = ragged_factory_ops.constant_value(
pylist, dtype=dtype, ragged_rank=ragged_rank, inner_shape=inner_shape)
# If dtype was explicitly specified, check it.
@@ -156,15 +159,15 @@
# If ragged_rank was explicitly specified, check it.
if ragged_rank is not None:
- if isinstance(rt, ragged.RaggedTensorValue):
+ if isinstance(rt, ragged_tensor_value.RaggedTensorValue):
self.assertEqual(rt.ragged_rank, ragged_rank)
else:
self.assertEqual(0, ragged_rank)
# If inner_shape was explicitly specified, check it.
if inner_shape is not None:
- if isinstance(rt, ragged.RaggedTensorValue):
- self.assertEqual(rt.inner_values.shape[1:], inner_shape)
+ if isinstance(rt, ragged_tensor_value.RaggedTensorValue):
+ self.assertEqual(rt.flat_values.shape[1:], inner_shape)
else:
self.assertEqual(rt.shape, inner_shape)
@@ -172,7 +175,10 @@
self.assertEqual(tuple(rt.shape), expected_shape)
if rt.shape:
- self.assertEqual(rt.tolist(), pylist)
+ if isinstance(rt, ragged_tensor_value.RaggedTensorValue):
+ self.assertEqual(rt.to_list(), pylist)
+ else:
+ self.assertEqual(rt.tolist(), pylist)
if expected_shape is not None:
self.assertEqual(rt.shape, expected_shape)
else:
@@ -252,11 +258,11 @@
inner_shape=None,
exception=None,
message=None):
- """Tests that `ragged.constant_value()` raises an expected exception."""
+ """Tests that `constant_value()` raises an expected exception."""
self.assertRaisesRegexp(
exception,
message,
- ragged.constant_value,
+ ragged_factory_ops.constant_value,
pylist,
dtype=dtype,
ragged_rank=ragged_rank,
diff --git a/tensorflow/python/ops/ragged/ragged_conversion_ops.py b/tensorflow/python/ops/ragged/ragged_conversion_ops.py
index 83212e4..854c5b3 100644
--- a/tensorflow/python/ops/ragged/ragged_conversion_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_conversion_ops.py
@@ -18,407 +18,27 @@
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 sparse_tensor
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import gen_ragged_conversion_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
-from tensorflow.python.ops.ragged import ragged_util
-#===============================================================================
-# RaggedTensor <-> Tensor conversion
-#===============================================================================
def from_tensor(tensor, lengths=None, padding=None, ragged_rank=1, name=None):
- """Converts a `Tensor` into a `RaggedTensor`.
-
- The set of absent/default values may be specified using a vector of lengths
- or a padding value (but not both). If `lengths` is specified, then the
- output tensor will satisfy `output[row] = tensor[row][:lengths[row]]`.
- If `padding` is specified, then any row *suffix* consisting entirely of
- `padding` will be excluded from the returned `RaggedTensor`. If neither
- `lengths` nor `padding` is specified, then the returned `RaggedTensor` will
- have no absent/default values.
-
- Examples:
-
- ```python
- >>> dt = tf.constant([[5, 7, 0], [0, 3, 0], [6, 0, 0]])
- >>> ragged.from_tensor(dt).eval().tolist()
- [[5, 7, 0], [0, 3, 0], [6, 0, 0]]
- >>> ragged.from_tensor(dt, lengths=[2, 0, 3]).eval().tolist()
- [[5, 7], [], [6, 0, 0]]
- >>> ragged.from_tensor(dt, padding=0).eval().tolist()
- [[5, 7], [0, 3], [6]]
- ```
-
- Args:
- tensor: The `Tensor` to convert. Must have rank `ragged_rank + 1` or
- higher.
- lengths: An optional set of row lengths, specified using a 1-D integer
- `Tensor` whose length is equal to `tensor.shape[0]` (the number of rows in
- `tensor`). If specified, then `output[row]` will contain
- `tensor[row][:lengths[row]]`. Negative lengths are treated as zero.
- padding: An optional padding value. If specified, then any row suffix
- consisting entirely of `padding` will be excluded from the returned
- RaggedTensor. `padding` is a `Tensor` with the same dtype as `tensor`
- and with `shape=tensor.shape[ragged_rank + 1:]`.
- ragged_rank: Integer specifying the ragged rank for the returned
- `RaggedTensor`. Must be greater than zero.
- name: A name prefix for the returned tensors (optional).
-
- Returns:
- A `RaggedTensor` with the specified `ragged_rank`. The shape of the
- returned ragged tensor is compatible with the shape of `tensor`.
- Raises:
- ValueError: If both `lengths` and `padding` are specified.
- """
- if lengths is not None and padding is not None:
- raise ValueError('Specify lengths or padding, but not both')
- if not isinstance(ragged_rank, int):
- raise TypeError('ragged_rank expected int, got %r' % ragged_rank)
- if ragged_rank <= 0:
- raise ValueError('ragged_rank must be greater than 0; got %s' % ragged_rank)
-
- with ops.name_scope(name, 'RaggedFromTensor', [tensor, lengths, padding]):
- tensor = ops.convert_to_tensor(tensor, name='tensor')
- tensor.shape.with_rank_at_least(ragged_rank + 1)
- input_shape = array_ops.shape(tensor, out_type=dtypes.int64)
- ncols = input_shape[1]
-
- # Handle ragged_rank>1 via recursion:
- # If the output should have multiple ragged dimensions, then first
- # flatten the tensor to eliminate all but the last ragged dimension,
- # and recursively convert that flattened tensor. Then add on the splits
- # for the dimensions that we flattened out.
- if ragged_rank > 1:
- # Flatten `tensor` to eliminate all but the last ragged dimension.
- new_shape = array_ops.concat(
- [constant_op.constant([-1], dtypes.int64), input_shape[ragged_rank:]],
- axis=0)
- flattened = array_ops.reshape(tensor, new_shape)
- # Recursively convert the flattened tensor.
- values = from_tensor(flattened, lengths, padding)
- # The total number of elements in each dimension. E.g., if
- # input_shape=[3, 4, 5, 6], then dim[2] has 3*4*5 elements in total.
- dim_size = math_ops.cumprod(input_shape)
- # Construct splits tensors for the dimensions that were flattened.
- new_splits = [
- math_ops.range(0, dim_size[dim - 1] + 1) * input_shape[dim]
- for dim in range(1, ragged_rank)
- ]
- return ragged_factory_ops.from_nested_row_splits(values, new_splits)
-
- # If padding was specified, then use it to find row lengths.
- if padding is not None:
- padding = ops.convert_to_tensor(
- padding, name='padding', dtype=tensor.dtype)
- padding.shape.assert_is_compatible_with(tensor.shape[2:])
-
- # Find places where the padding is equal to the tensor. (This will
- # broadcast `padding` across the outermost 2 dimensions of `tensor`,
- # so `has_default_value.shape = tensor.shape`.)
- has_default_value = math_ops.equal(padding, tensor)
-
- # If the padding isn't a scalar, then require that all values in the
- # padding match each item in the tensor. After this block of code,
- # `has_default.shape = tensor.shape[:2]`. (Unfortunately, we can't just
- # use reduce_all for both cases, becaue when you pass an empty `axis`
- # list to reduce_all, it reduces all axes; but we want it to reduce no
- # axes -- i.e., to be a no-op.)
- tensor_rank = array_ops.rank(tensor)
- reduce_axis = math_ops.range(2, tensor_rank)
- has_default = control_flow_ops.cond(
- tensor_rank > 2,
- lambda: math_ops.reduce_all(has_default_value, axis=reduce_axis),
- lambda: has_default_value)
- has_default.set_shape(tensor_shape.TensorShape([None, None]))
- has_default.set_shape(tensor.shape[:2])
-
- # Use has_default it to find the length of each row: for each non-default
- # item in a row, calculate the length that the row needs to have to
- # include that item; and then take the max of those values (across each
- # row).
- has_nondefault = math_ops.logical_not(has_default)
- has_nondefault = math_ops.cast(has_nondefault, dtypes.int64)
- length_for_nondefault_value = (
- has_nondefault * array_ops.expand_dims(
- math_ops.range(1, ncols + 1), 0))
- lengths = math_ops.reduce_max(length_for_nondefault_value, axis=1)
-
- # If we have lengths (either directly supplied, or computed from paddings),
- # then use those to construct splits; and then use masking to get the
- # corresponding values.
- if lengths is not None:
- lengths = ragged_util.convert_to_int_tensor(lengths, 'lengths',
- dtypes.int64)
- lengths.shape.assert_has_rank(1)
- lengths = math_ops.minimum(lengths, ncols)
- lengths = math_ops.maximum(lengths, 0)
- limits = math_ops.cumsum(lengths)
- splits = array_ops.concat(
- [array_ops.zeros([1], dtypes.int64), limits], axis=0)
- mask = array_ops.sequence_mask(lengths, maxlen=ncols)
- values = array_ops.boolean_mask(tensor, mask)
- return ragged_factory_ops.from_row_splits(values, splits)
-
- # If neither padding nor lengths were specified, then create a splits
- # vector that contains no default values, and reshape the input tensor
- # to form the values for the RaggedTensor.
- nrows = input_shape[0]
- nvals = nrows * ncols
- splits = math_ops.range(nrows + 1) * ncols
- values_shape = array_ops.concat([[nvals], input_shape[2:]], axis=0)
- values = array_ops.reshape(tensor, values_shape)
- return ragged_factory_ops.from_row_splits(values, splits)
+ if ragged_tensor.is_ragged(tensor):
+ return tensor
+ else:
+ return ragged_tensor.RaggedTensor.from_tensor(tensor, lengths, padding,
+ ragged_rank, name)
def to_tensor(rt_input, default_value=None, name=None):
- """Converts a `RaggedTensor` into a `Tensor`.
-
- Example:
-
- ```python
- >>> rt = ragged.constant([[9, 8, 7], [], [6, 5], [4]])
- >>> print ragged.to_tensor(rt).eval()
- [[9 8 7]
- [0 0 0]
- [6 5 0]
- [4 0 0]]
- ```
-
- Args:
- rt_input: The input `RaggedTensor`.
- default_value: Value to set for indices not specified in `rt_input`.
- Defaults to zero. `default_value` must be broadcastable to
- `rt_input.shape[rt_input.ragged_rank + 1:]`.
- name: A name prefix for the returned tensors (optional).
-
- Returns:
- A `Tensor` with shape `ragged.bounding_shape(rt_input)` and the
- values specified by the non-empty values in `rt_input`. Empty values are
- assigned `default_value`.
- """
- with ops.name_scope(name, 'RaggedToTensor', [rt_input, default_value]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
- if not ragged_tensor.is_ragged(rt_input):
- return rt_input # already dense
- if default_value is not None:
- default_value = ops.convert_to_tensor(
- default_value, name='default_value', dtype=rt_input.dtype)
-
- # If ragged_rank > 1, then recursively convert the ragged values into a
- # `Tensor` before we proceed.
- values = rt_input.values
- if ragged_tensor.is_ragged(values):
- values = to_tensor(values, default_value)
-
- # Tile the default value, if necessary.
- if default_value is not None:
- if values.shape.ndims is not None:
- default_value.shape.with_rank_at_most(values.shape.ndims - 1)
- if (values.shape.ndims is None or default_value.shape.ndims is None or
- values.shape.ndims != default_value.shape.ndims + 1):
- value_shape = array_ops.shape(values)[1:]
- default_value = array_ops.broadcast_to(default_value, value_shape)
- default_value.shape.assert_is_compatible_with(values.shape[1:])
-
- # Get the expected dense shape ([nrows, ncols] + value_shape).
- rt_row_lengths = [rt_input.row_splits[1:] - rt_input.row_splits[:-1]]
- nrows = array_ops.shape(rt_input.row_splits, out_type=dtypes.int64)[0] - 1
- ncols = math_ops.maximum(math_ops.reduce_max(rt_row_lengths), 0)
- values_shape = array_ops.shape(values, out_type=dtypes.int64)
- value_shape = values_shape[1:]
- nvals = values_shape[0]
-
- # Build a default value if none was supplied.
- if default_value is None:
- default_value = array_ops.zeros(value_shape, dtype=values.dtype)
- default_value.shape.assert_is_compatible_with(values.shape[1:])
- default_value.set_shape(values.shape[1:])
-
- # Get the row start indices, and expand to shape=[nrows, 1].
- starts = array_ops.expand_dims(rt_input.row_splits[:-1], 1)
-
- # Get the row limit indices, and expand to shape=[nrows, 1].
- limits = array_ops.expand_dims(rt_input.row_splits[1:], 1)
-
- # Get the column indices, and expand to shape=[1, ncols].
- columns = array_ops.expand_dims(math_ops.range(0, ncols), 0)
-
- # Build a list containing the values plus the default value. We will use
- # tf.gather to collect values from this list for the `Tensor` (using
- # nvals as the index for the default value).
- values_and_default = array_ops.concat(
- [values, array_ops.stack([default_value])], axis=0)
-
- # Construct a matrix "indices" pointing into values_and_default. I.e.,
- # output[r, c] = values_and_default[indices[r, c].
- nondefault_index = starts + columns
- has_value = nondefault_index < limits
- default_index = array_ops.fill(array_ops.stack([nrows, ncols]), nvals)
- indices = array_ops.where(has_value, nondefault_index, default_index)
-
- # Gather the results into a `Tensor`.
- return array_ops.gather(values_and_default, indices)
+ if ragged_tensor.is_ragged(rt_input):
+ return rt_input.to_tensor(default_value, name)
+ else:
+ return rt_input
-#===============================================================================
-# RaggedTensor <-> SparseTensor conversion
-#===============================================================================
def to_sparse(rt_input, name=None):
- """Converts a `RaggedTensor` into a sparse tensor.
-
- Example:
-
- ```python
- >>> rt = ragged.constant([[1, 2, 3], [4], [], [5, 6]])
- >>> ragged.to_sparse(rt).eval()
- SparseTensorValue(indices=[[0, 0], [0, 1], [0, 2], [1, 0], [3, 0], [3, 1]],
- values=[1, 2, 3, 4, 5, 6],
- dense_shape=[4, 3])
- ```
-
- Args:
- rt_input: The input `RaggedTensor`.
- name: A name prefix for the returned tensors (optional).
-
- Returns:
- A SparseTensor with the same values as `rt_input`.
- """
- if not ragged_tensor.is_ragged(rt_input):
- raise TypeError('Expected RaggedTensor, got %s' % type(rt_input).__name__)
- with ops.name_scope(name, 'RaggedToSparse', [rt_input]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
- result = gen_ragged_conversion_ops.ragged_tensor_to_sparse(
- rt_input.nested_row_splits, rt_input.inner_values, name=name)
- return sparse_tensor.SparseTensor(
- result.sparse_indices, result.sparse_values, result.sparse_dense_shape)
-
-
-@ops.RegisterGradient('RaggedTensorToSparse')
-def _ragged_tensor_to_sparse_gradient(op, unused_sparse_indices_grad,
- sparse_values_grad,
- unused_sparse_shape_grad):
- """Gradient for ragged.to_sparse."""
- op_inputs_nested_row_splits = op.inputs[:-1]
- op_inputs_inner_values = op.inputs[-1]
-
- # No gradient for the RaggedTensor's nested_row_splits.
- nested_row_splits_gradient = [None] * len(op_inputs_nested_row_splits)
-
- # Gradient for the RaggedTensor's inner_values is formed by reshaping
- # the gradient for the SparseTensor's values.
- inner_values_shape = array_ops.shape(op_inputs_inner_values)
- inner_values_gradient = array_ops.reshape(sparse_values_grad,
- inner_values_shape)
-
- return nested_row_splits_gradient + [inner_values_gradient]
+ return rt_input.to_sparse(name)
def from_sparse(st_input, name=None):
- """Converts a 2D `SparseTensor` to a `RaggedTensor`.
-
- Each row of the `output` `RaggedTensor` will contain the explicit values from
- the same row in `st_input`. `st_input` must be ragged-right. If not it is
- not ragged-right, then an error will be generated.
-
- Example:
-
- ```python
- >>> st = SparseTensor(indices=[[0, 1], [0, 2], [0, 3], [1, 0], [3, 0]],
- ... values=[1, 2, 3, 4, 5],
- ... dense_shape=[4, 3])
- >>> ragged.from_sparse(st).eval().tolist()
- [[1, 2, 3], [4], [], [5]]
- ```
-
- Currently, only two-dimensional `SparseTensors` are supported.
-
- Args:
- st_input: The sparse tensor to convert. Must have rank 2.
- name: A name prefix for the returned tensors (optional).
-
- Returns:
- A `RaggedTensor` with the same values as `st_input`.
- `output.ragged_rank = rank(st_input) - 1`.
- `output.shape = [st_input.dense_shape[0], None]`.
- Raises:
- ValueError: If the number of dimensions in `st_input` is not known
- statically, or is not two.
- """
- if not sparse_tensor.is_sparse(st_input):
- raise TypeError('Expected SparseTensor, got %s' % type(st_input).__name__)
- with ops.name_scope(name, 'RaggedFromSparse', [st_input]):
- st_input = sparse_tensor.convert_to_tensor_or_sparse_tensor(
- st_input, name='rt_input')
-
- static_rank_from_dense_shape = (
- None if st_input.dense_shape.shape.ndims is None
- else st_input.dense_shape.shape.dims[0].value)
- static_rank_from_indices = (
- None if st_input.indices.shape.ndims is None
- else st_input.indices.shape.dims[1].value)
-
- if static_rank_from_dense_shape != 2 and static_rank_from_indices != 2:
- raise ValueError('rank(st_input) must be 2')
-
- with ops.control_dependencies(
- _assert_sparse_indices_are_ragged_right(st_input.indices)):
- # Treat sparse row indices as segment ids to generate a splits tensor that
- # we can pair with the sparse tensor values. (Ignore sparse column
- # indices.)
- segment_ids = st_input.indices[:, 0]
- num_segments = st_input.dense_shape[0]
- return ragged_factory_ops.from_value_rowids(st_input.values, segment_ids,
- num_segments)
-
-
-def _assert_sparse_indices_are_ragged_right(indices):
- """Checks that the given SparseTensor.indices tensor is ragged-right.
-
- Example: `indices = [[0, 0], [0, 1], [2, 0], [3, 1]]` is not ragged right
- because the entry `[3, 1]` skips a cell.
-
- Args:
- indices: The SparseTensor indices to check.
-
- Returns:
- A list of control dependency op tensors.
- """
- index_prefix = indices[:, :-1]
- index_suffix = indices[:, -1]
-
- # Check whether each index is starting a new row in the innermost dimension
- # (prefix[i] != prefix[i-1]) or continuing a row (prefix[i] == prefix[i-1]).
- # (Note: this skips the first index; we will check that separately below.)
- index_prefix_changed = math_ops.reduce_any(
- math_ops.not_equal(index_prefix[1:], index_prefix[:-1]), axis=1)
-
- # Check two cases:
- # * For indices that start a new row: index_suffix[i] must be zero.
- # * For indices that continue a row: index_suffix[i] must be equal to
- # index_suffix[i-1]+1.
- index_ok = array_ops.where(
- index_prefix_changed, math_ops.equal(index_suffix[1:], 0),
- math_ops.equal(index_suffix[1:], index_suffix[:-1] + 1))
-
- # Also check that the very first index didn't skip any cells. The first
- # index starts a new row (by definition), so its suffix should be zero.
- sparse_indices_are_ragged_right = math_ops.logical_and(
- math_ops.reduce_all(math_ops.equal(index_suffix[:1], 0)),
- math_ops.reduce_all(index_ok))
-
- message = [
- 'SparseTensor is not right-ragged',
- 'SparseTensor.indices =', indices
- ]
- return [control_flow_ops.Assert(sparse_indices_are_ragged_right, message)]
+ return ragged_tensor.RaggedTensor.from_sparse(st_input, name)
diff --git a/tensorflow/python/ops/ragged/ragged_dispatch.py b/tensorflow/python/ops/ragged/ragged_dispatch.py
new file mode 100644
index 0000000..ecc7f5d
--- /dev/null
+++ b/tensorflow/python/ops/ragged/ragged_dispatch.py
@@ -0,0 +1,465 @@
+# 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.
+# ==============================================================================
+"""Operator dispatch for RaggedTensors."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.util import dispatch
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_export
+from tensorflow.python.util import tf_inspect
+
+# @TODO(edloper): Set this to True in the CL that exports RaggedTensors.
+_UPDATE_DOCSTRINGS = False
+
+# Information about an argument to an operation: The name of the argument, its
+# position in the argument list, and a boolean flag indicating whether it
+# expects a list of tensors.
+_ArgInfo = collections.namedtuple('ArgInfo', ['name', 'position', 'is_list'])
+
+
+def _get_arg_infos(func, arg_names):
+ """Returns an `_ArgInfo` for each argument of `func` specified by `arg_names`.
+
+ Args:
+ func: The function whose arguments should be described.
+ arg_names: The names of the arguments to get info for.
+
+ Returns:
+ A tuple of `_ArgInfo`s.
+ """
+ arg_infos = []
+
+ # Inspect the func's argspec to find the position of each arg.
+ arg_spec = tf_inspect.getargspec(func)
+ for argname in arg_names:
+ assert isinstance(argname, str)
+ is_list = argname.startswith('[') and argname.endswith(']')
+ if is_list:
+ argname = argname[1:-1]
+ if argname not in arg_spec.args:
+ raise ValueError('Argument %r not found function in %s. Args=%s' %
+ (argname, func, arg_spec.args))
+ arg_infos.append(_ArgInfo(argname, arg_spec.args.index(argname), is_list))
+ return arg_infos
+
+
+def _is_convertible_to_tensor(value):
+ """Returns true if `value` is convertible to a `Tensor`."""
+ if value is None:
+ return True
+ if isinstance(value,
+ (ops.Tensor, variables.Variable, np.ndarray, int, float, str)):
+ return True
+ elif isinstance(value, (sparse_tensor.SparseTensor,)):
+ return False
+ else:
+ try:
+ ops.convert_to_tensor(value)
+ return True
+ except (TypeError, ValueError):
+ return False
+
+
+class UnaryRaggedElementwiseDispatcher(dispatch.OpDispatcher):
+ """OpDispatcher for unary ops that map a base op across ragged values."""
+
+ def __init__(self, original_op, arg_is_list=False):
+ self._original_op = original_op
+ self._arg_is_list = arg_is_list
+ arg_names = tf_inspect.getfullargspec(original_op)[0]
+ self._x = arg_names[0]
+ if _UPDATE_DOCSTRINGS:
+ original_op.__doc__ = (
+ original_op.__doc__.rstrip() + '\n\n' +
+ ' `{x}` may be a `tf.RaggedTensor`.\n'.format(x=self._x))
+
+ def handle(self, args, kwargs):
+ if args:
+ x, args = args[0], args[1:]
+ else:
+ kwargs = kwargs.copy()
+ x = kwargs.pop(self._x, None)
+ if x is None:
+ return self.NOT_SUPPORTED
+ if self._arg_is_list:
+ found_ragged = False
+ for elt in x:
+ if ragged_tensor.is_ragged(elt):
+ found_ragged = True
+ elif not _is_convertible_to_tensor(elt):
+ return self.NOT_SUPPORTED
+ if found_ragged:
+ nested_splits_lists = [
+ elt.nested_row_splits for elt in x if ragged_tensor.is_ragged(elt)
+ ]
+ flat_values = [
+ elt.flat_values if ragged_tensor.is_ragged(elt) else elt
+ for elt in x
+ ]
+ with ops.control_dependencies(
+ ragged_util.assert_splits_match(nested_splits_lists)):
+ return ragged_tensor.RaggedTensor.from_nested_row_splits(
+ self._original_op(flat_values, *args, **kwargs),
+ nested_splits_lists[0])
+ else:
+ return self.NOT_SUPPORTED
+ else:
+ found_ragged = ragged_tensor.is_ragged(x)
+ if found_ragged:
+ mapped_values = self._original_op(x.flat_values, *args, **kwargs)
+ return x.with_flat_values(mapped_values)
+ else:
+ return self.NOT_SUPPORTED
+
+
+class BinaryRaggedElementwiseDispatcher(dispatch.OpDispatcher):
+ """OpDispatcher for binary ops that map a base op across ragged values.
+
+ Supports broadcasting.
+ """
+
+ def __init__(self, original_op):
+ self._original_op = original_op
+ arg_names = tf_inspect.getfullargspec(original_op)[0]
+ self._x = arg_names[0]
+ self._y = arg_names[1]
+ if _UPDATE_DOCSTRINGS:
+ original_op.__doc__ = (
+ original_op.__doc__.rstrip() + '\n\n' +
+ ' `{x}` and `{y}` may be a `tf.RaggedTensor`.\n'.format(
+ x=self._x, y=self._y))
+
+ def handle(self, args, kwargs):
+ # Extract the binary args.
+ if len(args) > 1:
+ x = args[0]
+ y = args[1]
+ args = args[2:]
+ elif args:
+ kwargs = kwargs.copy()
+ x = args[0]
+ y = kwargs.pop(self._y, None)
+ args = args[1:]
+ else:
+ kwargs = kwargs.copy()
+ x = kwargs.pop(self._x, None)
+ y = kwargs.pop(self._y, None)
+
+ # Bail if we don't have at least one ragged argument.
+ x_is_ragged = ragged_tensor.is_ragged(x)
+ y_is_ragged = ragged_tensor.is_ragged(y)
+ if not (x_is_ragged or y_is_ragged):
+ return self.NOT_SUPPORTED
+
+ # Convert args to tensors. Bail if conversion fails.
+ try:
+ if not x_is_ragged:
+ x = ops.convert_to_tensor(x, name=self._x, preferred_dtype=y.dtype)
+ if not y_is_ragged:
+ y = ops.convert_to_tensor(y, name=self._y, preferred_dtype=x.dtype)
+ except (TypeError, ValueError):
+ return self.NOT_SUPPORTED
+
+ if ((x_is_ragged and y_is_ragged) or
+ (x_is_ragged and x.flat_values.shape.ndims <= y.shape.ndims) or
+ (y_is_ragged and y.flat_values.shape.ndims <= x.shape.ndims)):
+ bcast_shape = ragged_tensor_shape.broadcast_dynamic_shape(
+ ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(x),
+ ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(y))
+ x = ragged_tensor_shape.broadcast_to(
+ x, bcast_shape, broadcast_inner_dimensions=False)
+ y = ragged_tensor_shape.broadcast_to(
+ y, bcast_shape, broadcast_inner_dimensions=False)
+
+ x_values = x.flat_values if ragged_tensor.is_ragged(x) else x
+ y_values = y.flat_values if ragged_tensor.is_ragged(y) else y
+ mapped_values = self._original_op(x_values, y_values, *args, **kwargs)
+ if ragged_tensor.is_ragged(x):
+ return x.with_flat_values(mapped_values)
+ else:
+ return y.with_flat_values(mapped_values)
+
+
+class RaggedDispatcher(dispatch.OpDispatcher):
+ """OpDispatcher for ragged ops.
+
+ Dispatches to a wrapped op-handler if at least one of the `tensor_args`
+ arguments is a RaggedTensor or a RaggedTensorValue; and all of the
+ `tensor_args` arguments are convertible to Tensor or RaggedTensor.
+ """
+
+ def __init__(self, original_op, ragged_op, ragged_args):
+ op_arg_names = tf_inspect.getfullargspec(original_op)[0]
+ ragged_arg_names = tf_inspect.getfullargspec(ragged_op)[0]
+ if op_arg_names != ragged_arg_names:
+ raise AssertionError(
+ 'Signature must exactly match when overriding %s with %s: %s vs %s' %
+ (original_op, ragged_op, op_arg_names, ragged_arg_names))
+ self._ragged_op = ragged_op
+ self._ragged_args = _get_arg_infos(ragged_op, ragged_args)
+ if _UPDATE_DOCSTRINGS:
+ arg_list = ' and '.join('`%s`' % arg for arg in ragged_args)
+ original_op.__doc__ = (
+ original_op.__doc__.rstrip() + '\n\n' +
+ ' {0} may be a `tf.RaggedTensor`.\n'.format(arg_list))
+
+ def handle(self, args, kwargs):
+ if self.is_supported(args, kwargs):
+ return self._ragged_op(*args, **kwargs)
+ else:
+ return self.NOT_SUPPORTED
+
+ def is_supported(self, args, kwargs):
+ found_ragged = False
+ for arg_info in self._ragged_args:
+ if arg_info.position < len(args):
+ arg = args[arg_info.position]
+ else:
+ arg = kwargs.get(arg_info.name, None)
+
+ if arg_info.is_list:
+ if not isinstance(arg, (list, tuple)):
+ return False
+ for elt in arg:
+ if ragged_tensor.is_ragged(elt):
+ found_ragged = True
+ elif not _is_convertible_to_tensor(elt):
+ return False
+ else:
+ if ragged_tensor.is_ragged(arg):
+ found_ragged = True
+ elif not _is_convertible_to_tensor(arg):
+ return False
+ return found_ragged
+
+
+def ragged_dispatch(original_op, tensor_args):
+
+ def decorator(ragged_op):
+ dispatch.RaggedDispatcher(original_op, ragged_op,
+ tensor_args).register(original_op)
+ return ragged_op
+
+ return decorator
+
+
+_UNARY_ELEMENTWISE_OPS = [
+ array_ops.check_numerics,
+ array_ops.identity,
+ array_ops.ones_like,
+ array_ops.ones_like_v2,
+ array_ops.zeros_like,
+ array_ops.zeros_like_v2,
+ clip_ops.clip_by_value,
+ math_ops.abs,
+ math_ops.acos,
+ math_ops.acosh,
+ math_ops.angle,
+ math_ops.asin,
+ math_ops.asinh,
+ math_ops.atan,
+ math_ops.atanh,
+ math_ops.cast,
+ math_ops.ceil,
+ math_ops.conj,
+ math_ops.cos,
+ math_ops.cosh,
+ math_ops.digamma,
+ math_ops.erf,
+ math_ops.erfc,
+ math_ops.exp,
+ math_ops.expm1,
+ math_ops.floor,
+ math_ops.imag,
+ math_ops.is_finite,
+ math_ops.is_inf,
+ math_ops.is_nan,
+ math_ops.lgamma,
+ math_ops.log,
+ math_ops.log1p,
+ math_ops.log_sigmoid,
+ math_ops.logical_not,
+ math_ops.negative,
+ math_ops.real,
+ math_ops.reciprocal,
+ math_ops.rint,
+ math_ops.round,
+ math_ops.rsqrt,
+ math_ops.saturate_cast,
+ math_ops.sign,
+ math_ops.sin,
+ math_ops.sinh,
+ math_ops.sqrt,
+ math_ops.square,
+ math_ops.tan,
+ parsing_ops.decode_compressed,
+ string_ops.string_to_number,
+ string_ops.string_to_hash_bucket,
+ string_ops.as_string,
+ string_ops.decode_base64,
+ string_ops.encode_base64,
+ string_ops.regex_full_match,
+ string_ops.regex_replace,
+ string_ops.string_strip,
+ string_ops.string_to_hash_bucket,
+ string_ops.string_to_hash_bucket_fast,
+ string_ops.string_to_hash_bucket_strong,
+ string_ops.substr,
+ string_ops.substr_v2,
+ string_ops.string_length,
+ string_ops.string_length_v2,
+ string_ops.unicode_script,
+]
+
+_UNARY_LIST_ELEMENTWISE_OPS = [
+ math_ops.add_n,
+ string_ops.string_join,
+]
+
+_BINARY_ELEMENTWISE_OPS = [
+ math_ops.add,
+ math_ops.atan2,
+ math_ops.complex,
+ math_ops.div_no_nan,
+ math_ops.divide,
+ math_ops.equal,
+ math_ops.floordiv,
+ math_ops.floormod,
+ math_ops.greater,
+ math_ops.greater_equal,
+ math_ops.less,
+ math_ops.less_equal,
+ math_ops.logical_and,
+ math_ops.logical_or,
+ math_ops.logical_xor,
+ math_ops.maximum,
+ math_ops.minimum,
+ math_ops.multiply,
+ math_ops.not_equal,
+ math_ops.pow,
+ math_ops.realdiv,
+ math_ops.squared_difference,
+ math_ops.subtract,
+ math_ops.truediv,
+ math_ops.truncatediv,
+ math_ops.truncatemod,
+]
+
+
+def _ragged_gather_v1(params, indices, validate_indices=None, name=None,
+ axis=0):
+ return ragged_array_ops.gather(params=params, indices=indices,
+ validate_indices=validate_indices,
+ axis=axis, name=name)
+
+
+def _ragged_expand_dims_v1(input, axis=None, name=None, dim=None): # pylint: disable=redefined-builtin
+ if dim is not None:
+ axis = dim
+ return ragged_array_ops.expand_dims(input=input, axis=axis, name=name)
+
+
+# (original_op, ragged_op, ragged_args)
+_RAGGED_DISPATCH_OPS = [
+ (array_ops.batch_gather, ragged_array_ops.batch_gather,
+ ['params', 'indices']),
+ (array_ops.concat, ragged_array_ops.concat, ['[values]']),
+ (array_ops.expand_dims, _ragged_expand_dims_v1, ['input']),
+ (array_ops.expand_dims_v2, ragged_array_ops.expand_dims, ['input']),
+ (array_ops.gather, _ragged_gather_v1, ['params', 'indices']),
+ (array_ops.gather_v2, ragged_array_ops.gather, ['params', 'indices']),
+ (array_ops.gather_nd, ragged_array_ops.gather_nd, ['params', 'indices']),
+ (array_ops.stack, ragged_array_ops.stack, ['[values]']),
+ (array_ops.tile, ragged_array_ops.tile, ['input']),
+ (array_ops.where, ragged_array_ops.where, ['condition', 'x', 'y']),
+ (math_ops.unsorted_segment_sum, ragged_math_ops.segment_sum,
+ ['data', 'segment_ids']),
+ (math_ops.unsorted_segment_prod, ragged_math_ops.segment_prod,
+ ['data', 'segment_ids']),
+ (math_ops.unsorted_segment_min, ragged_math_ops.segment_min,
+ ['data', 'segment_ids']),
+ (math_ops.unsorted_segment_max, ragged_math_ops.segment_max,
+ ['data', 'segment_ids']),
+ (math_ops.unsorted_segment_mean, ragged_math_ops.segment_mean,
+ ['data', 'segment_ids']),
+ (math_ops.unsorted_segment_sqrt_n, ragged_math_ops.segment_sqrt_n,
+ ['data', 'segment_ids']),
+ (math_ops.reduce_sum, ragged_math_ops.reduce_sum, ['input_tensor']),
+ (math_ops.reduce_prod, ragged_math_ops.reduce_prod, ['input_tensor']),
+ (math_ops.reduce_min, ragged_math_ops.reduce_min, ['input_tensor']),
+ (math_ops.reduce_max, ragged_math_ops.reduce_max, ['input_tensor']),
+ (math_ops.reduce_mean, ragged_math_ops.reduce_mean, ['input_tensor']),
+ (math_ops.reduce_any, ragged_math_ops.reduce_any, ['input_tensor']),
+ (math_ops.reduce_all, ragged_math_ops.reduce_all, ['input_tensor']),
+]
+
+
+def register_dispatchers():
+ """Constructs & registers OpDispatchers for ragged ops."""
+
+ op_list = (
+ _UNARY_ELEMENTWISE_OPS + _UNARY_LIST_ELEMENTWISE_OPS +
+ _BINARY_ELEMENTWISE_OPS + [x[0] for x in _RAGGED_DISPATCH_OPS])
+ for op in op_list:
+ _, undecorated_op = tf_decorator.unwrap(op)
+ if not hasattr(undecorated_op, tf_export.API_ATTRS['tensorflow'].names):
+ raise AssertionError('Expected %s to be an exported symbol '
+ '(while adding a RaggedTensor dispatcher)')
+
+ for op in _UNARY_ELEMENTWISE_OPS:
+ UnaryRaggedElementwiseDispatcher(op).register(op)
+
+ for op in _UNARY_LIST_ELEMENTWISE_OPS:
+ UnaryRaggedElementwiseDispatcher(op, True).register(op)
+
+ for op in _BINARY_ELEMENTWISE_OPS:
+ BinaryRaggedElementwiseDispatcher(op).register(op)
+
+ for (original_op, ragged_op, args) in _RAGGED_DISPATCH_OPS:
+ RaggedDispatcher(original_op, ragged_op, args).register(original_op)
+
+
+def ragged_op_list():
+ """Returns a string listing operators that have dispathers registered."""
+ op_list = (
+ _UNARY_ELEMENTWISE_OPS + _UNARY_LIST_ELEMENTWISE_OPS +
+ _BINARY_ELEMENTWISE_OPS + [x[0] for x in _RAGGED_DISPATCH_OPS])
+ return (
+ '\n\n### Additional ops that support `RaggedTensor`\n\n' + '\n'.join([
+ '* `tf.%s`' % tf_export.get_canonical_name_for_symbol(op)
+ for op in op_list
+ ]))
+
+
+register_dispatchers()
diff --git a/tensorflow/python/ops/ragged/ragged_dispatch_test.py b/tensorflow/python/ops/ragged/ragged_dispatch_test.py
new file mode 100644
index 0000000..9d70470
--- /dev/null
+++ b/tensorflow/python/ops/ragged/ragged_dispatch_test.py
@@ -0,0 +1,687 @@
+# 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 RaggedTensor operator dispatch."""
+
+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.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.platform import googletest
+
+# Constants listing various op types to test. Each operation
+# should be included in at least one list below, or tested separately if
+# necessary (e.g., because it expects additional arguments).
+UNARY_FLOAT_OPS = [
+ math_ops.abs,
+ math_ops.acos,
+ math_ops.acosh,
+ math_ops.angle,
+ math_ops.asin,
+ math_ops.asinh,
+ math_ops.atan,
+ math_ops.atanh,
+ math_ops.ceil,
+ math_ops.conj,
+ math_ops.cos,
+ math_ops.cosh,
+ math_ops.digamma,
+ math_ops.erf,
+ math_ops.erfc,
+ math_ops.exp,
+ math_ops.expm1,
+ math_ops.floor,
+ math_ops.imag,
+ math_ops.is_finite,
+ math_ops.is_inf,
+ math_ops.is_nan,
+ math_ops.lgamma,
+ math_ops.log,
+ math_ops.log1p,
+ math_ops.log_sigmoid,
+ math_ops.negative,
+ math_ops.real,
+ math_ops.reciprocal,
+ math_ops.rint,
+ math_ops.round,
+ math_ops.rsqrt,
+ math_ops.sign,
+ math_ops.sin,
+ math_ops.sinh,
+ math_ops.sqrt,
+ math_ops.square,
+ math_ops.tan,
+ array_ops.identity,
+ array_ops.ones_like,
+ array_ops.zeros_like,
+]
+UNARY_BOOL_OPS = [
+ math_ops.logical_not,
+]
+UNARY_STRING_OPS = [
+ string_ops.decode_base64,
+ string_ops.encode_base64,
+ string_ops.string_strip,
+ parsing_ops.decode_compressed,
+]
+BINARY_FLOAT_OPS = [
+ math_ops.add,
+ math_ops.atan2,
+ math_ops.complex,
+ math_ops.div_no_nan,
+ math_ops.divide,
+ math_ops.equal,
+ math_ops.floordiv,
+ math_ops.floormod,
+ math_ops.greater,
+ math_ops.greater_equal,
+ math_ops.less,
+ math_ops.less_equal,
+ math_ops.maximum,
+ math_ops.minimum,
+ math_ops.multiply,
+ math_ops.not_equal,
+ math_ops.pow,
+ math_ops.realdiv,
+ math_ops.squared_difference,
+ math_ops.subtract,
+ math_ops.truediv,
+]
+BINARY_BOOL_OPS = [
+ math_ops.logical_and,
+ math_ops.logical_or,
+ math_ops.logical_xor,
+]
+UNARY_INT_OPS = [
+ string_ops.unicode_script,
+]
+BINARY_INT_OPS = [
+ math_ops.truncatediv,
+ math_ops.truncatemod,
+]
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedElementwiseOpsTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
+
+ def assertSameShape(self, x, y):
+ """Checks that x and y have the same shape (including ragged shapes)."""
+ if isinstance(x, ragged_tensor.RaggedTensor):
+ self.assertIsInstance(y, ragged_tensor.RaggedTensor)
+ self.assertEqual(x.ragged_rank, y.ragged_rank)
+ for (x_splits, y_splits) in zip(x.nested_row_splits, y.nested_row_splits):
+ self.assertAllEqual(x_splits, y_splits)
+ self.assertAllEqual(
+ array_ops.shape(x.flat_values), array_ops.shape(y.flat_values))
+ else:
+ self.assertIsInstance(y, ops.Tensor)
+ self.assertAllEqual(array_ops.shape(x), array_ops.shape(y))
+
+ @parameterized.parameters(
+ #=========================================================================
+ # Test different input shapes.
+ #=========================================================================
+ [
+ # 0-dimensional input
+ {'x': 12},
+ # 1-dimensional input
+ {'x': [1, -2, 3]},
+ # 2-dimensional input
+ {'x': [[-2, 3], [-3, 4]]},
+ {'x': ragged_factory_ops.constant_value(
+ [[-2, 3], [-3]], ragged_rank=1)},
+ # 3-dimensional inputs
+ {'x': [[[-2, 3], [3, 4]], [[7, 6], [5, 4]]]},
+ {'x': ragged_factory_ops.constant_value(
+ [[[-2, 3], [3, 4]], [[7, 6]]],
+ ragged_rank=1)},
+ {'x': ragged_factory_ops.constant_value(
+ [[[-2, 3, 4], []], [[7, 6]], []],
+ ragged_rank=2)},
+ ] +
+ #=========================================================================
+ # Test each unary op.
+ #=========================================================================
+ [{'x': ragged_factory_ops.constant_value([[-2.0, 3.0], [-3.0]]), 'op': op}
+ for op in UNARY_FLOAT_OPS] +
+ [{'x': ragged_factory_ops.constant_value([[True, False], [True]]),
+ 'op': op}
+ for op in UNARY_BOOL_OPS] +
+ [{'x': ragged_factory_ops.constant_value([[18, 512], [12412]], np.int32),
+ 'op': op}
+ for op in UNARY_INT_OPS] +
+ [{'x': ragged_factory_ops.constant_value([['abcd', 'efgh'],
+ ['aabbccdd']]),
+ 'op': op}
+ for op in UNARY_STRING_OPS] +
+ [
+ {'op': clip_ops.clip_by_value,
+ 'x': ragged_factory_ops.constant_value([[-2.0, 3.0], [-3.0]]),
+ 'clip_value_min': 0.1, 'clip_value_max': 4.0},
+ {'op': math_ops.cast,
+ 'x': ragged_factory_ops.constant_value([[-2.0, 3.0], [-3.0]]),
+ 'dtype': dtypes.int32},
+ {'op': math_ops.saturate_cast,
+ 'x': ragged_factory_ops.constant_value([[-2.0, 3.0], [-3.0]]),
+ 'dtype': dtypes.int32},
+ {'op': string_ops.string_to_hash_bucket,
+ 'x': ragged_factory_ops.constant_value(
+ [['abcd', 'efgh'], ['aabbccdd']]),
+ 'num_buckets': 1000},
+ {'op': string_ops.string_to_hash_bucket_fast,
+ 'x': ragged_factory_ops.constant_value(
+ [['abcd', 'efgh'], ['aabbccdd']]),
+ 'num_buckets': 1000},
+ {'op': string_ops.string_to_hash_bucket_strong,
+ 'x': ragged_factory_ops.constant_value(
+ [['abcd', 'efgh'], ['aabbccdd']]),
+ 'num_buckets': 1000,
+ 'key': [1231, 12512]},
+ {'op': string_ops.string_to_number,
+ 'x': ragged_factory_ops.constant_value([['-2.0', '3.0'], ['-3.0']])},
+ {'op': string_ops.regex_full_match,
+ 'x': ragged_factory_ops.constant_value([['hello', '123'], ['1+1']]),
+ 'pattern': r'\w+'},
+ {'op': string_ops.regex_replace,
+ 'x': ragged_factory_ops.constant_value([['hello', '123'], ['1+1']]),
+ 'pattern': r'\d',
+ 'rewrite': '#'},
+ {'op': string_ops.substr,
+ 'x': ragged_factory_ops.constant_value([['hello', '123'], ['1+1']]),
+ 'pos': 2, 'len': 3},
+ {'op': array_ops.check_numerics,
+ 'x': ragged_factory_ops.constant_value([[-2.0, 3.0], [-3.0]]),
+ 'message': 'check-numerics'},
+ ]
+ ) # pyformat: disable
+ def testUnaryElementwiseOp(self, x, op=math_ops.abs, **extra_args):
+ x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x)
+ result = op(x, **extra_args)
+
+ # Run the wrapped op on the dense values, for comparison.
+ dense_x = x.flat_values if isinstance(x, ragged_tensor.RaggedTensor) else x
+ expected_flat_values = array_ops.reshape(op(dense_x, **extra_args), [-1])
+
+ # Check that the result has the expected shape.
+ self.assertSameShape(x, result)
+
+ # Check that the result has the expected (flattened) values.
+ if isinstance(result, ragged_tensor.RaggedTensor):
+ result_flat_values = array_ops.reshape(result.flat_values, [-1])
+ else:
+ result_flat_values = array_ops.reshape(result, [-1])
+ self.assertAllEqual(expected_flat_values, result_flat_values)
+
+ @parameterized.parameters(
+ [
+ #=====================================================================
+ # Without broadcasting -- i.e., shapes match exactly.
+ #=====================================================================
+ # Shapes: x:(), y:()
+ {'x': 12,
+ 'y': 8},
+ # Shapes: x:(3,), y:(3,)
+ {'x': [7, 8, 9],
+ 'y': [1, -2, 3]},
+ # Shapes: x:(2, 2), y:(2, 2)
+ {'x': [[-2, 3], [-3, -4]],
+ 'y': [[1, 2], [3, 4]]},
+ # Shapes: x:(2, None), y:(2, None)
+ {'x': ragged_factory_ops.constant_value([[-2, 3], [-3]]),
+ 'y': ragged_factory_ops.constant_value([[5, 6], [7]])},
+ # Shapes: x:(2, 2, 2), y:(2, 2, 2)
+ {'x': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+ 'y': [[[9, 3], [3, 4]], [[5, 2], [7, 6]]]},
+ # Shapes: x:(2, None, None), y: (2, None, None)
+ {'x': ragged_factory_ops.constant_value(
+ [[[1, 2], [3], [4]], [[], [5, 7, 8]]]),
+ 'y': ragged_factory_ops.constant_value(
+ [[[3, 8], [2], [5]], [[], [1, 9, 8]]])},
+ # Shapes: x:(2, None, 2), y: (2, None, 2)
+ {'x': ragged_factory_ops.constant_value(
+ [[[1, 2]], [[3, 4], [5, 6], [7, 8]]],
+ ragged_rank=1),
+ 'y': ragged_factory_ops.constant_value(
+ [[[9, 3]], [[5, 2], [3, 4], [7, 6]]],
+ ragged_rank=1)},
+
+ #=====================================================================
+ # With broadcasting
+ #=====================================================================
+ # Shapes: x:(), y:(3,)
+ {'x': 12, # Broadcast () -> (3,)
+ 'y': [1, -2, 3]},
+ # Shapes: x:(1,), y:(3,)
+ {'x': [12], # Broadcast (1,) -> (3,)
+ 'y': [1, -2, 3]},
+ # Shapes: x:(), y:(2, 2)
+ {'x': 12, # Broadcast () -> (2, 2)
+ 'y': [[1, 2], [3, 4]]},
+ # Shapes: x:(1,), y:(2, 2)
+ {'x': 12, # Broadcast (1,) -> (2, 2)
+ 'y': [[1, 2], [3, 4]]},
+ # Shapes: x:(2, 1), y:(2, 2)
+ {'x': [[10], [20]], # Broadcast (2, 1) -> (2, 2)
+ 'y': [[1, 2], [3, 4]]},
+ # Shapes: x:(), y:(2, None)
+ {'x': 10, # Broadcast () -> (2, None)
+ 'y': ragged_factory_ops.constant_value(
+ [[1, 2], [3]], dtype=np.int32)},
+ # TODO(edloper): Add tests for more advanced broadcasting, once we add
+ # support for it.
+
+ #=====================================================================
+ # Keyword Args
+ #=====================================================================
+ {'x': ragged_factory_ops.constant_value(
+ [[[1, 2], [3], [4]], [[], [5, 7, 8]]]),
+ 'y': ragged_factory_ops.constant_value(
+ [[[3, 8], [2], [5]], [[], [1, 9, 8]]]),
+ 'use_kwargs': ('x', 'y')},
+ {'x': ragged_factory_ops.constant_value(
+ [[[1, 2]], [[3, 4], [5, 6], [7, 8]]],
+ ragged_rank=1),
+ 'y': ragged_factory_ops.constant_value(
+ [[[9, 3]], [[5, 2], [3, 4], [7, 6]]],
+ ragged_rank=1),
+ 'use_kwargs': ('x', 'y')},
+ {'x': ragged_factory_ops.constant_value(
+ [[[1, 2]], [[3, 4], [5, 6], [7, 8]]],
+ ragged_rank=1),
+ 'y': ragged_factory_ops.constant_value(
+ [[[9, 3]], [[5, 2], [3, 4], [7, 6]]],
+ ragged_rank=1),
+ 'use_kwargs': ('x',)},
+ ] +
+ #=========================================================================
+ # Test each unary op.
+ #=========================================================================
+ [{'x': ragged_factory_ops.constant_value([[-2.0, 3.0], [-3.0]]),
+ 'y': ragged_factory_ops.constant_value([[5.0, 1.0], [12.0]]),
+ 'op': op}
+ for op in BINARY_FLOAT_OPS] +
+ [{'x': ragged_factory_ops.constant_value([[-2, 3], [-3]]),
+ 'y': ragged_factory_ops.constant_value([[5, 1], [12]]),
+ 'op': op}
+ for op in BINARY_INT_OPS] +
+ [{'x': ragged_factory_ops.constant_value([[True, True], [False]]),
+ 'y': ragged_factory_ops.constant_value([[False, True], [False]]),
+ 'op': op}
+ for op in BINARY_BOOL_OPS]
+ ) # pyformat: disable
+ def testBinaryElementwiseOp(self, x, y, op=math_ops.add, **extra_args):
+ use_kwargs = extra_args.pop('use_kwargs', ())
+ x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x)
+ y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y)
+ if 'x' in use_kwargs and 'y' in use_kwargs:
+ result = op(x=x, y=y, **extra_args)
+ elif 'y' in use_kwargs:
+ result = op(x, y=y, **extra_args)
+ else:
+ result = op(x, y, **extra_args)
+
+ # Run the wrapped op on the dense values, for comparison.
+ dense_x = x.flat_values if isinstance(x, ragged_tensor.RaggedTensor) else x
+ dense_y = y.flat_values if isinstance(y, ragged_tensor.RaggedTensor) else y
+ expected_flat_values = array_ops.reshape(
+ op(dense_x, dense_y, **extra_args), [-1])
+
+ # Check that the result has the expected shape.
+ self.assertSameShape(y, result)
+
+ # Check that the result has the expected (flattened) values.
+ if isinstance(result, ragged_tensor.RaggedTensor):
+ result_flat_values = array_ops.reshape(result.flat_values, [-1])
+ else:
+ result_flat_values = array_ops.reshape(result, [-1])
+ self.assertAllEqual(expected_flat_values, result_flat_values)
+
+ @parameterized.parameters(
+ [
+ {'inputs': (12, 8, 3)},
+ {'inputs': ([1, 2, 3], [7, 8, 9], [3, 6, 9])},
+ {'inputs': ([[1, 2]], [[3, 4]], [[5, 6]])},
+ {'inputs': (ragged_factory_ops.constant_value([[1, 3], [-3]]),
+ ragged_factory_ops.constant_value([[4, 7], [88]]),
+ ragged_factory_ops.constant_value([[2, 9], [12]]))},
+ {'inputs': (ragged_factory_ops.constant_value(
+ [[[1, 3], [-3]], [[1]]]),
+ ragged_factory_ops.constant_value(
+ [[[4, 7], [88]], [[2]]]),
+ ragged_factory_ops.constant_value(
+ [[[2, 9], [12]], [[8]]]))},
+ {'inputs': (
+ ragged_factory_ops.constant_value([[[1, 3], [3, 4]], [[1, 5]]],
+ ragged_rank=1),
+ ragged_factory_ops.constant_value([[[4, 7], [1, 2]], [[2, 2]]],
+ ragged_rank=1),
+ ragged_factory_ops.constant_value([[[2, 9], [5, 2]], [[8, 0]]],
+ ragged_rank=1))},
+ {'inputs': (
+ ragged_factory_ops.constant_value([[[1, 3], [-3]], [[1]]]),
+ ragged_factory_ops.constant_value([[[4, 7], [88]], [[2]]]),
+ ragged_factory_ops.constant_value([[[2, 9], [12]], [[8]]])),
+ 'use_kwargs': True},
+ ] + [
+ {'op': math_ops.add_n,
+ 'inputs': (ragged_factory_ops.constant_value([[1, 3], [-3]]),
+ ragged_factory_ops.constant_value([[4, 7], [88]]),
+ ragged_factory_ops.constant_value([[2, 9], [12]]))},
+ {'op': string_ops.string_join,
+ 'inputs': (
+ ragged_factory_ops.constant_value([['a', 'b'], ['c']]),
+ ragged_factory_ops.constant_value([['foo', 'bar'], ['baz']]),
+ ragged_factory_ops.constant_value([['2', '9'], ['12']]))},
+ ]) # pyformat: disable
+ def testListValuedElementwiseOp(self, inputs, op=math_ops.add_n,
+ **extra_args):
+ use_kwargs = extra_args.pop('use_kwargs', False)
+ inputs = [
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(x) for x in inputs
+ ]
+ if use_kwargs:
+ result = op(inputs=inputs, **extra_args)
+ else:
+ result = op(inputs, **extra_args)
+
+ # Run the wrapped op on the dense values, for comparison.
+ dense_inputs = [
+ x.flat_values if isinstance(x, ragged_tensor.RaggedTensor) else x
+ for x in inputs
+ ]
+ expected_flat_values = array_ops.reshape(
+ op(dense_inputs, **extra_args), [-1])
+
+ # Check that the result has the expected shape.
+ self.assertSameShape(inputs[0], result)
+
+ # Check that the result has the expected (flattened) values.
+ if isinstance(result, ragged_tensor.RaggedTensor):
+ result_flat_values = array_ops.reshape(result.flat_values, [-1])
+ else:
+ result_flat_values = array_ops.reshape(result, [-1])
+ self.assertAllEqual(expected_flat_values, result_flat_values)
+
+ def testElementwiseOpUnknownRankError(self):
+ if context.executing_eagerly():
+ return
+ x = ragged_factory_ops.constant([[1, 2], [3]])
+ y = ragged_tensor.RaggedTensor.from_row_splits(
+ array_ops.placeholder_with_default([1, 2, 3], shape=None), x.row_splits)
+ with self.assertRaisesRegexp(ValueError,
+ r'Unable to broadcast: unknown rank'):
+ math_ops.add(x, y)
+
+ @parameterized.parameters([
+ dict(
+ x=ragged_factory_ops.constant_value([[1, 2], [3]]),
+ y=[[10]],
+ expected=[[11, 12], [13]]),
+ dict(
+ x=ragged_factory_ops.constant_value([[[1, 2], [3, 4]], [[5]]],
+ ragged_rank=2),
+ y=ragged_factory_ops.constant_value([[[10], [20]], [[30]]],
+ ragged_rank=1),
+ expected=[[[11, 12], [23, 24]], [[35]]]),
+ dict(
+ x=ragged_factory_ops.constant_value([[[1]]]),
+ y=ragged_factory_ops.constant_value([[1]]),
+ expected=[[[2]]]),
+ ])
+ def testElementwiseOpBroadcast(self, x, y, expected):
+ x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x, dtype=dtypes.int32)
+ y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y, dtype=dtypes.int32)
+ result = x + y
+ self.assertRaggedEqual(result, expected)
+
+ def testElementwiseOpShapeMismatch(self):
+ x = ragged_factory_ops.constant([[1, 2, 3], [4, 5]])
+ y = ragged_factory_ops.constant([[1, 2, 3], [4, 5, 6]])
+ with self.assertRaises(errors.InvalidArgumentError):
+ self.evaluate(math_ops.add(x, y))
+
+ def testBinaryOpSparseAndRagged(self):
+ x = ragged_factory_ops.constant([[1, 2, 3], [4, 5]])
+ y = sparse_tensor.SparseTensor([[0, 0], [0, 1], [2, 0]], [1, 2, 3], [3, 2])
+ with self.assertRaises((TypeError, ValueError)):
+ self.evaluate(math_ops.add(x, y))
+
+ with self.assertRaises((TypeError, ValueError)):
+ self.evaluate(math_ops.add_n([x, y]))
+
+ @parameterized.parameters([
+ dict(
+ op=array_ops.batch_gather,
+ args=(ragged_factory_ops.constant_value([[5, 6, 7], [8, 9]]),
+ ragged_factory_ops.constant_value([[2, 1, 0], [1]])),
+ expected=ragged_factory_ops.constant_value([[7, 6, 5], [9]])),
+ dict(
+ op=array_ops.concat,
+ args=([
+ ragged_factory_ops.constant_value([[1, 2, 3], [4]],
+ dtype=np.int32),
+ np.array([[5, 6]], dtype=np.int32)
+ ],),
+ kwargs={'axis': 0},
+ expected=ragged_factory_ops.constant_value([[1, 2, 3], [4], [5, 6]])),
+ dict(
+ op=array_ops.expand_dims,
+ kwargs={
+ 'input': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'axis': 0
+ },
+ expected=ragged_factory_ops.constant_value([[[1, 2], [3]]])),
+ dict(
+ op=array_ops.expand_dims_v2,
+ kwargs={
+ 'input': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'axis': -1
+ },
+ expected=ragged_factory_ops.constant_value([[[1], [2]], [[3]]],
+ ragged_rank=1),
+ ),
+ dict(
+ op=array_ops.gather,
+ kwargs={
+ 'params': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'indices': [1, 0, 1]
+ },
+ expected=ragged_factory_ops.constant_value([[3], [1, 2], [3]])),
+ dict(
+ op=array_ops.gather_v2,
+ kwargs={
+ 'params': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'indices': ragged_factory_ops.constant_value([[1, 0], [1]])
+ },
+ expected=ragged_factory_ops.constant_value([[[3], [1, 2]], [[3]]])),
+ dict(
+ op=array_ops.gather_nd,
+ kwargs={
+ 'params': ragged_factory_ops.constant_value([[7, 8], [9]]),
+ 'indices': [[0, 1], [1, 0], [0, 0]]
+ },
+ expected=ragged_factory_ops.constant_value([8, 9, 7])),
+ dict(
+ op=array_ops.stack,
+ args=([
+ ragged_factory_ops.constant_value([[1, 2, 3], [4]],
+ dtype=np.int32),
+ np.array([[5, 6]], dtype=np.int32)
+ ],),
+ expected=ragged_factory_ops.constant_value([[[1, 2, 3], [4]],
+ [[5, 6]]])),
+ dict(
+ op=array_ops.tile,
+ args=([
+ ragged_factory_ops.constant_value([[1, 2], [3]], dtype=np.int32),
+ [2, 3]
+ ]),
+ expected=ragged_factory_ops.constant_value([[1, 2, 1, 2, 1, 2],
+ [3, 3, 3],
+ [1, 2, 1, 2, 1, 2],
+ [3, 3, 3]])),
+ dict(
+ op=array_ops.where,
+ args=(ragged_factory_ops.constant_value([[True, False], [True]]),
+ ragged_factory_ops.constant_value([[b'A', b'B'], [b'C']]),
+ ragged_factory_ops.constant_value([[b'a', b'b'], [b'c']])),
+ expected=ragged_factory_ops.constant_value([[b'A', b'b'], [b'C']])),
+ dict(
+ op=array_ops.where,
+ args=(ragged_factory_ops.constant_value([[True, False], [True]]),),
+ expected=[[0, 0], [1, 0]]),
+ dict(
+ op=math_ops.unsorted_segment_sum,
+ kwargs={
+ 'data': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'segment_ids': ragged_factory_ops.constant_value([[0, 2], [0]]),
+ 'num_segments': 3
+ },
+ expected=[4, 0, 2]),
+ dict(
+ op=math_ops.unsorted_segment_prod,
+ kwargs={
+ 'data': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'segment_ids': ragged_factory_ops.constant_value([[0, 2], [0]]),
+ 'num_segments': 3
+ },
+ expected=[3, 1, 2]),
+ dict(
+ op=math_ops.unsorted_segment_min,
+ kwargs={
+ 'data': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'segment_ids': ragged_factory_ops.constant_value([[0, 1], [0]]),
+ 'num_segments': 2
+ },
+ expected=[1, 2]),
+ dict(
+ op=math_ops.unsorted_segment_max,
+ kwargs={
+ 'data': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'segment_ids': ragged_factory_ops.constant_value([[0, 1], [0]]),
+ 'num_segments': 2
+ },
+ expected=[3, 2]),
+ dict(
+ op=math_ops.unsorted_segment_mean,
+ kwargs={
+ 'data': ragged_factory_ops.constant_value([[1, 2], [3]]),
+ 'segment_ids': ragged_factory_ops.constant_value([[0, 1], [0]]),
+ 'num_segments': 2
+ },
+ expected=[2, 2]),
+ dict(
+ op=math_ops.unsorted_segment_sqrt_n,
+ kwargs={
+ 'data':
+ ragged_factory_ops.constant_value([[1.0, 2.0],
+ [3.0, 4.0, 6.0]]),
+ 'segment_ids':
+ ragged_factory_ops.constant_value([[0, 1], [0, 0, 0]]),
+ 'num_segments':
+ 2
+ },
+ expected=[7.0, 2.0]),
+ dict(
+ op=math_ops.reduce_sum,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[1, 2], [3, 4, 5]]),
+ 'axis':
+ 1
+ },
+ expected=[3, 12]),
+ dict(
+ op=math_ops.reduce_prod,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[1, 2], [3, 4, 5]]),
+ 'axis':
+ 1
+ },
+ expected=[2, 60]),
+ dict(
+ op=math_ops.reduce_min,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[1, 2], [3, 4, 5]]),
+ 'axis':
+ 1
+ },
+ expected=[1, 3]),
+ dict(
+ op=math_ops.reduce_max,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[1, 2], [3, 4, 5]]),
+ 'axis':
+ 1
+ },
+ expected=[2, 5]),
+ dict(
+ op=math_ops.reduce_mean,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[1, 3], [3, 4, 5]]),
+ 'axis':
+ 1
+ },
+ expected=[2, 4]),
+ dict(
+ op=math_ops.reduce_any,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[True, False],
+ [True, True, True]]),
+ 'axis':
+ 1
+ },
+ expected=[True, True]),
+ dict(
+ op=math_ops.reduce_all,
+ kwargs={
+ 'input_tensor':
+ ragged_factory_ops.constant_value([[True, False],
+ [True, True, True]]),
+ 'axis':
+ 1
+ },
+ expected=[False, True]),
+ ])
+ def testRaggedDispatch(self, op, expected, args=(), kwargs=None):
+ if kwargs is None: kwargs = {}
+ result = op(*args, **kwargs)
+ self.assertRaggedEqual(result, expected)
+
+
+if __name__ == '__main__':
+ googletest.main()
diff --git a/tensorflow/python/ops/ragged/ragged_eager_test.py b/tensorflow/python/ops/ragged/ragged_eager_test.py
index 731ff74..86f01aa 100644
--- a/tensorflow/python/ops/ragged/ragged_eager_test.py
+++ b/tensorflow/python/ops/ragged/ragged_eager_test.py
@@ -17,17 +17,17 @@
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-import sys
from absl.testing import parameterized
from tensorflow.python.framework import ops
-from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+class RaggedTensorTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters([
dict(pylist=[[b'a', b'b'], [b'c']]),
@@ -35,22 +35,16 @@
dict(pylist=[[[1, 2], [3, 4]], [[5, 6], [], [7, 8]]], ragged_rank=1),
])
def testRaggedTensorToList(self, pylist, ragged_rank=None):
- rt = ragged.constant(pylist, ragged_rank)
- self.assertEqual(rt.tolist(), pylist)
-
- expected = "RaggedTensor([['a', 'b'], ['c']])"
- if sys.version_info[0] == 3:
- expected = "RaggedTensor([[b'a', b'b'], [b'c']])"
+ rt = ragged_factory_ops.constant(pylist, ragged_rank)
+ self.assertRaggedEqual(rt, pylist)
@parameterized.parameters([
- dict(pylist=[['a', 'b'], ['c']],
- expected=expected),
- dict(pylist=[[[1, 2], [3]], [[4, 5, 6], [], [7]]],
- expected='RaggedTensor([[[1, 2], [3]], [[4, 5, 6], [], [7]]])'),
+ dict(pylist=[[b'a', b'b'], [b'c']]),
+ dict(pylist=[[[1, 2], [3]], [[4, 5, 6], [], [7]]]),
])
- def testRaggedTensorStr(self, pylist, expected):
- rt = ragged.constant(pylist)
- self.assertEqual(str(rt), expected)
+ def testRaggedTensorStr(self, pylist):
+ rt = ragged_factory_ops.constant(pylist)
+ self.assertEqual(str(rt), '<tf.RaggedTensor %s>' % pylist)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_elementwise_ops.py b/tensorflow/python/ops/ragged/ragged_elementwise_ops.py
deleted file mode 100644
index 59b7dd1..0000000
--- a/tensorflow/python/ops/ragged/ragged_elementwise_ops.py
+++ /dev/null
@@ -1,389 +0,0 @@
-# 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.
-# ==============================================================================
-"""Elementwise operations for RaggedTensors."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import collections
-
-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 math_ops
-from tensorflow.python.ops import parsing_ops
-from tensorflow.python.ops import string_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
-from tensorflow.python.ops.ragged import ragged_tensor
-from tensorflow.python.ops.ragged import ragged_tensor_shape
-from tensorflow.python.util import tf_decorator
-from tensorflow.python.util import tf_export
-from tensorflow.python.util import tf_inspect
-
-# Information about an argument to an operation: The name of the argument, its
-# position in the argument list, and a boolean flag indicating whether it
-# expects a list of tensors.
-_ArgInfo = collections.namedtuple('ArgInfo', ['name', 'position', 'is_list'])
-
-
-def make_elementwise_op(op, *elementwise_args):
- """Returns a ragged-tensor version of the elementwise operation `op`.
-
- The returned operation will:
-
- 1. Broadcast the elementwise arguments to have a compatible shape.
- An exception is raised if the tensors not broadcast-compatible.
- 2. Call `op`, substituting the dense values of the broadcasted tensor for
- each elementwise argument.
- 3. Return a potentially ragged tensor constructed from the output of `op`
- and the broadcasted tensors' nested row splits.
-
- For example, you can construct a ragged-tensor version of the standard
- operation `tf.add` by calling `make_elementwise_op(tf.add, 'x', 'y')`.
-
- Args:
- op: The operation to wrap.
- *elementwise_args: The names of arguments to `op` that are treated as
- elementwise. Arguments that take a list of tensors should have their
- names wrapped in square brackets (e.g. "[inputs]").
-
- Raises:
- ValueError: If any name specified in `elementwise_args` is not the name
- of an argument to `op`.
- """
- elementwise_arg_infos = _get_arg_infos(op, elementwise_args)
-
- def ragged_op(*args, **kwargs):
- """Ragged version of `op`."""
- args = list(args)
-
- # Collect all of the elementwise arguments, and put them in a single
- # dict whose values are the (potentially ragged) tensors that need to
- # be broadcast to a common shape. The keys of this dict are tuples
- # (argkey, index), where argkey is an int for poitional args or a string
- # for keyword args; and index is None for non-list args and the index of the
- # tensor for list args.
- elementwise_args = {}
- for (name, position, is_list) in elementwise_arg_infos.values():
- if position < len(args):
- if is_list:
- args[position] = list(args[position])
- for (index, arg) in enumerate(args[position]):
- elementwise_args[position, index] = arg
- else:
- elementwise_args[position, None] = args[position]
- elif name in kwargs:
- if is_list:
- kwargs[name] = list(kwargs[name])
- for (i, arg) in enumerate(kwargs[name]):
- elementwise_args[name, i] = arg
- else:
- elementwise_args[name, None] = kwargs[name]
-
- with ops.name_scope(None, op.__name__, elementwise_args.values()):
- # Convert all inputs to tensors or ragged tensors.
- for ((key, index), tensor) in elementwise_args.items():
- argname = elementwise_arg_infos[key].name
- converted = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- tensor, name=argname)
- elementwise_args[key, index] = converted
-
- # Broadcast tensors to have compatible shapes.
- broadcast_args, result_splits, broadcast_check_ops = \
- _broadcast_elementwise_args(elementwise_args)
-
- # Replace tensor arguments with their dense values.
- for ((key, index), tensor) in broadcast_args.items():
- if ragged_tensor.is_ragged(tensor):
- if isinstance(key, int) and index is None:
- args[key] = tensor.inner_values
- elif isinstance(key, int) and index is not None:
- args[key][index] = tensor.inner_values
- elif isinstance(key, str) and index is None:
- kwargs[key] = tensor.inner_values
- else:
- assert isinstance(key, str) and index is not None
- kwargs[key][index] = tensor.inner_values
-
- # Call the elementwise op on the broadcasted dense values.
- with ops.control_dependencies(broadcast_check_ops):
- result_values = op(*args, **kwargs)
-
- # Restore any ragged dimensions that we stripped off, and return the
- # result.
- return ragged_factory_ops.from_nested_row_splits(result_values,
- result_splits)
-
- # Construct the docstring.
- op_name = tf_export.get_canonical_name_for_symbol(op)
- assert op_name is not None, op
- argnames = ', '.join('`%s`' % s.strip('[]') for s in elementwise_args)
- docstring = _ELEMENTWISE_DOCSTRING % dict(op_name=op_name, argnames=argnames)
-
- # Update name, docstring, signature, etc., for the wrapper, and return it.
- return tf_decorator.make_decorator(op, ragged_op, decorator_doc=docstring)
-
-
-_ELEMENTWISE_DOCSTRING = """\
-Ragged version of the elementwise operation `tf.%(op_name)s`.
-
- The following elementwise arguments may be ragged or dense:
- %(argnames)s.
- These arguments will be broadcast to a compatible shape if necessary.
- """
-
-
-def _get_arg_infos(func, elementwise_args):
- """Returns `_ArgInfo`s for each `func` arg specified by `elementwise_args`.
-
- Args:
- func: The function whose arguments should be described.
- elementwise_args: The names of the arguments to get info for.
-
- Returns:
- A dictionary that maps both names and positions of arguments to
- `_ArgInfo` tuples.
- """
- arg_infos = {}
-
- # Inspect the func's argspec to find the position of each arg.
- arg_spec = tf_inspect.getargspec(func)
- for argname in elementwise_args:
- assert isinstance(argname, str)
- is_list = argname.startswith('[') and argname.endswith(']')
- if is_list:
- argname = argname[1:-1]
- assert argname in arg_spec.args, (func, argname, arg_spec.args)
- arg_info = _ArgInfo(argname, arg_spec.args.index(argname), is_list)
- arg_infos[arg_info.name] = arg_info
- arg_infos[arg_info.position] = arg_info
- return arg_infos
-
-
-def _broadcast_elementwise_args(elementwise_args):
- """Broadcasts the values of `elementwise_args` to have compatible shapes.
-
- Args:
- elementwise_args: A dictionary whose keys are potentially ragged tensors.
-
- Returns:
- A tuple `(broadcast_args, broadcast_splits, checks)` where:
-
- * `broadcast_args` is a dictionary with the same keys as
- `elementwise_args`, mapping to broadcasted tensors.
- * `broadcast_splits` is the broadcasted nested row splits.
- * `checks` is a possibly empty tuple of assertion operations that should
- be added as control dependencies.
-
- Raises:
- ValueError: If broadcasting fails.
- """
- # No elementwise arguments were used: nothing to do!
- if not elementwise_args:
- return elementwise_args, (), ()
-
- # A single elementwise argument was used: no broadcasting necessary.
- if len(elementwise_args) == 1:
- arg = list(elementwise_args.values())[0]
- if ragged_tensor.is_ragged(arg):
- return elementwise_args, arg.nested_row_splits, ()
- else:
- return elementwise_args, (), ()
-
- # Multiple elementwise arguments.
- else:
- is_ragged = [ragged_tensor.is_ragged(t) for t in elementwise_args.values()]
- if not any(is_ragged):
- return elementwise_args, (), ()
-
- # If we have a single ragged tensor plus a set of scalars, then we can
- # rely on the underlying elementwise op to do broadcasting.
- if (sum(is_ragged) == 1 and
- all((ragged_tensor.is_ragged(t) or t.shape.ndims == 0)
- for t in elementwise_args.values())):
- nested_splits_lists = [
- t.nested_row_splits
- for t in elementwise_args.values()
- if ragged_tensor.is_ragged(t)][0]
- return elementwise_args, nested_splits_lists, ()
-
- else:
- # Get the shapes of all the elementwise arguments.
- shapes = [ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(t)
- for t in elementwise_args.values()]
-
- # Broadcast the shapes to all have the same rank (the max rank).
- ranks = [t.shape.ndims for t in elementwise_args.values()]
- if any(rank is None for rank in ranks):
- raise ValueError('Unable to broadcast: unknown rank')
- broadcast_rank = max(ranks)
- shapes = [shape.broadcast_to_rank(broadcast_rank) for shape in shapes]
-
- # For each dimension, broadcast the shapes to be compatible.
- for axis in range(broadcast_rank):
- # For each i, broadcast shape[i+1] to be compatible with shape[i]; and
- # then finally broadcast shape[0] to be compatible with shape[-1].
- for i in range(len(shapes)):
- j = (i + 1) % len(shapes)
- dim_size = shapes[i].dimension_size(axis)
- shapes[j] = shapes[j].broadcast_dimension(axis, dim_size)
- broadcast_shape = shapes[0]
-
- # Broadcast every elementwise arg to the shape that we calculated.
- elementwise_args = dict([
- (key, ragged_tensor_shape.broadcast_to(t, broadcast_shape, False))
- for (key, t) in elementwise_args.items()])
- nested_splits_lists = list(elementwise_args.values())[0].nested_row_splits
- return elementwise_args, nested_splits_lists, ()
-
-
-# A list of symbols that should be exported in the "ragged" package.
-_symbols_to_export = []
-
-
-def _add_elementwise_ops_to_this_module(specs, verbose=False):
- """Adds ragged versions of the given ops to this module.
-
- Args:
- specs: A list of tuples containing the arguments for `make_elementwise_op`.
- verbose: If true, then display each op that gets added.
- """
- for spec in specs:
- original_op = spec[0]
- ragged_op = make_elementwise_op(*spec)
- canonical_name = tf_export.get_canonical_name_for_symbol(original_op)
- if '.' not in canonical_name:
- op_name = canonical_name
- else:
- op_name = original_op.__name__
-
- # Temporary hack (will be removed once dispatch is added for RaggedTensors):
- if op_name == 'neg': op_name = 'negative'
-
- if verbose:
- print('Adding ragged_elementwise_op: tf.ragged.%s (based on tf.%s)' %
- (op_name, canonical_name))
- globals()[op_name] = ragged_op
- _symbols_to_export.append(op_name)
-
-
-# A list of tuples containing arguments for `make_elementwise_op`, for each
-# elementwise operation that should have a ragged version built. Each tuple
-# contains a standard `Tensor` operation, and the names of any arguments
-# that are processed in elementwise fashion.
-_TF_ELEMENTWISE_OPS = [
- # Unary math operations.
- (clip_ops.clip_by_value, 't'),
- (math_ops.abs, 'x'),
- (math_ops.acos, 'x'),
- (math_ops.acosh, 'x'),
- (math_ops.angle, 'input'),
- (math_ops.asin, 'x'),
- (math_ops.asinh, 'x'),
- (math_ops.atan, 'x'),
- (math_ops.atanh, 'x'),
- (math_ops.cast, 'x'),
- (math_ops.ceil, 'x'),
- (math_ops.conj, 'x'),
- (math_ops.cos, 'x'),
- (math_ops.cosh, 'x'),
- (math_ops.digamma, 'x'),
- (math_ops.erf, 'x'),
- (math_ops.erfc, 'x'),
- (math_ops.exp, 'x'),
- (math_ops.expm1, 'x'),
- (math_ops.floor, 'x'),
- (math_ops.imag, 'input'),
- (math_ops.is_finite, 'x'),
- (math_ops.is_inf, 'x'),
- (math_ops.is_nan, 'x'),
- (math_ops.lgamma, 'x'),
- (math_ops.log, 'x'),
- (math_ops.log1p, 'x'),
- (math_ops.log_sigmoid, 'x'),
- (math_ops.logical_not, 'x'),
- (math_ops.negative, 'x'),
- (math_ops.real, 'input'),
- (math_ops.reciprocal, 'x'),
- (math_ops.rint, 'x'),
- (math_ops.round, 'x'),
- (math_ops.rsqrt, 'x'),
- (math_ops.saturate_cast, 'value'),
- (math_ops.sign, 'x'),
- (math_ops.sin, 'x'),
- (math_ops.sinh, 'x'),
- (math_ops.sqrt, 'x'),
- (math_ops.square, 'x'),
- (math_ops.tan, 'x'),
-
- # Binary math operations
- (math_ops.add, 'x', 'y'),
- (math_ops.atan2, 'y', 'x'),
- (math_ops.complex, 'real', 'imag'),
- (math_ops.div, 'x', 'y'),
- (math_ops.div_no_nan, 'x', 'y'),
- (math_ops.divide, 'x', 'y'),
- (math_ops.equal, 'x', 'y'),
- (math_ops.floordiv, 'x', 'y'),
- (math_ops.floormod, 'x', 'y'),
- (math_ops.greater, 'x', 'y'),
- (math_ops.greater_equal, 'x', 'y'),
- (math_ops.less, 'x', 'y'),
- (math_ops.less_equal, 'x', 'y'),
- (math_ops.logical_and, 'x', 'y'),
- (math_ops.logical_or, 'x', 'y'),
- (math_ops.logical_xor, 'x', 'y'),
- (math_ops.maximum, 'x', 'y'),
- (math_ops.minimum, 'x', 'y'),
- (math_ops.multiply, 'x', 'y'),
- (math_ops.not_equal, 'x', 'y'),
- (math_ops.pow, 'x', 'y'),
- (math_ops.realdiv, 'x', 'y'),
- (math_ops.squared_difference, 'x', 'y'),
- (math_ops.subtract, 'x', 'y'),
- (math_ops.truediv, 'x', 'y'),
- (math_ops.truncatediv, 'x', 'y'),
- (math_ops.truncatemod, 'x', 'y'),
-
- # N-ary math operations
- (math_ops.add_n, '[inputs]'),
-
- # String operations
- (string_ops.as_string, 'input'),
- (string_ops.decode_base64, 'input'),
- (string_ops.encode_base64, 'input'),
- (string_ops.regex_full_match, 'input'),
- (string_ops.regex_replace, 'input'),
- (string_ops.string_join, '[inputs]'),
- (string_ops.string_strip, 'input'),
- (string_ops.string_to_hash_bucket, 'input'),
- (string_ops.string_to_hash_bucket_fast, 'input'),
- (string_ops.string_to_hash_bucket_strong, 'input'),
- (string_ops.substr, 'input'),
- (string_ops.unicode_script, 'input'),
-
- # Array ops
- (array_ops.check_numerics, 'tensor'),
- (array_ops.identity, 'input'),
- (array_ops.ones_like, 'tensor'),
- (array_ops.zeros_like, 'tensor'),
-
- # Parsing ops
- (parsing_ops.decode_compressed, 'bytes'),
- (parsing_ops.string_to_number, 'string_tensor'),
-]
-_add_elementwise_ops_to_this_module(_TF_ELEMENTWISE_OPS)
-
diff --git a/tensorflow/python/ops/ragged/ragged_elementwise_ops_test.py b/tensorflow/python/ops/ragged/ragged_elementwise_ops_test.py
deleted file mode 100644
index 305a96d..0000000
--- a/tensorflow/python/ops/ragged/ragged_elementwise_ops_test.py
+++ /dev/null
@@ -1,443 +0,0 @@
-# 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 ragged.elementwise_ops."""
-
-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.framework import dtypes
-from tensorflow.python.framework import errors
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import test_util
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
-from tensorflow.python.platform import googletest
-
-# Constants listing various op types to test. Each elementwise operation
-# should be included in at least one list below, or tested separately if
-# necessary (e.g., because it expects additional arguments).
-UNARY_FLOAT_OPS = [
- ragged.abs,
- ragged.acos,
- ragged.acosh,
- ragged.angle,
- ragged.asin,
- ragged.asinh,
- ragged.atan,
- ragged.atanh,
- ragged.ceil,
- ragged.conj,
- ragged.cos,
- ragged.cosh,
- ragged.digamma,
- ragged.erf,
- ragged.erfc,
- ragged.exp,
- ragged.expm1,
- ragged.floor,
- ragged.imag,
- ragged.is_finite,
- ragged.is_inf,
- ragged.is_nan,
- ragged.lgamma,
- ragged.log,
- ragged.log1p,
- ragged.log_sigmoid,
- ragged.negative,
- ragged.real,
- ragged.reciprocal,
- ragged.rint,
- ragged.round,
- ragged.rsqrt,
- ragged.sign,
- ragged.sin,
- ragged.sinh,
- ragged.sqrt,
- ragged.square,
- ragged.tan,
- ragged.as_string,
- ragged.identity,
- ragged.ones_like,
- ragged.zeros_like,
-]
-UNARY_BOOL_OPS = [
- ragged.logical_not,
-]
-UNARY_STRING_OPS = [
- ragged.decode_base64,
- ragged.encode_base64,
- ragged.string_strip,
- ragged.decode_compressed,
-]
-BINARY_FLOAT_OPS = [
- ragged.add,
- ragged.atan2,
- ragged.complex,
- ragged.div,
- ragged.div_no_nan,
- ragged.divide,
- ragged.equal,
- ragged.floordiv,
- ragged.floormod,
- ragged.greater,
- ragged.greater_equal,
- ragged.less,
- ragged.less_equal,
- ragged.maximum,
- ragged.minimum,
- ragged.multiply,
- ragged.not_equal,
- ragged.pow,
- ragged.realdiv,
- ragged.squared_difference,
- ragged.subtract,
- ragged.truediv,
-]
-BINARY_BOOL_OPS = [
- ragged.logical_and,
- ragged.logical_or,
- ragged.logical_xor,
-]
-UNARY_INT_OPS = [
- ragged.unicode_script,
-]
-BINARY_INT_OPS = [
- ragged.truncatediv,
- ragged.truncatemod,
-]
-
-
-class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase,
- parameterized.TestCase):
-
- def assertSameShape(self, x, y):
- """Checks that x and y have the same shape (including ragged shapes)."""
- if isinstance(x, ragged.RaggedTensor):
- self.assertIsInstance(y, ragged.RaggedTensor)
- self.assertEqual(x.ragged_rank, y.ragged_rank)
- for (x_splits, y_splits) in zip(x.nested_row_splits, y.nested_row_splits):
- self.assertAllEqual(x_splits, y_splits)
- self.assertAllEqual(
- array_ops.shape(x.inner_values), array_ops.shape(y.inner_values))
- else:
- self.assertIsInstance(y, ops.Tensor)
- self.assertAllEqual(array_ops.shape(x), array_ops.shape(y))
-
- @parameterized.parameters(
- #=========================================================================
- # Test different input shapes.
- #=========================================================================
- [
- # 0-dimensional input
- {'x': 12},
- # 1-dimensional input
- {'x': [1, -2, 3]},
- # 2-dimensional input
- {'x': [[-2, 3], [-3, 4]]},
- {'x': ragged.constant_value([[-2, 3], [-3]], ragged_rank=1)},
- # 3-dimensional inputs
- {'x': [[[-2, 3], [3, 4]], [[7, 6], [5, 4]]]},
- {'x': ragged.constant_value([[[-2, 3], [3, 4]], [[7, 6]]],
- ragged_rank=1)},
- {'x': ragged.constant_value([[[-2, 3, 4], []], [[7, 6]], []],
- ragged_rank=2)},
- ] +
- #=========================================================================
- # Test each unary op.
- #=========================================================================
- [{'x': ragged.constant_value([[-2.0, 3.0], [-3.0]]), 'op': op}
- for op in UNARY_FLOAT_OPS] +
- [{'x': ragged.constant_value([[True, False], [True]]), 'op': op}
- for op in UNARY_BOOL_OPS] +
- [{'x': ragged.constant_value([[18, 512], [12412]], np.int32), 'op': op}
- for op in UNARY_INT_OPS] +
- [{'x': ragged.constant_value([['abcd', 'efgh'], ['aabbccdd']]), 'op': op}
- for op in UNARY_STRING_OPS] +
- [
- {'op': ragged.clip_by_value,
- 'x': ragged.constant_value([[-2.0, 3.0], [-3.0]]),
- 'clip_value_min': 0.1, 'clip_value_max': 4.0},
- {'op': ragged.cast,
- 'x': ragged.constant_value([[-2.0, 3.0], [-3.0]]),
- 'dtype': dtypes.int32},
- {'op': ragged.saturate_cast,
- 'x': ragged.constant_value([[-2.0, 3.0], [-3.0]]),
- 'dtype': dtypes.int32},
- {'op': ragged.string_to_hash_bucket,
- 'x': ragged.constant_value([['abcd', 'efgh'], ['aabbccdd']]),
- 'num_buckets': 1000},
- {'op': ragged.string_to_hash_bucket_fast,
- 'x': ragged.constant_value([['abcd', 'efgh'], ['aabbccdd']]),
- 'num_buckets': 1000},
- {'op': ragged.string_to_hash_bucket_strong,
- 'x': ragged.constant_value([['abcd', 'efgh'], ['aabbccdd']]),
- 'num_buckets': 1000,
- 'key': [1231, 12512]},
- {'op': ragged.string_to_number,
- 'x': ragged.constant_value([['-2.0', '3.0'], ['-3.0']])},
- {'op': ragged.regex_full_match,
- 'x': ragged.constant_value([['hello', '123'], ['1+1']]),
- 'pattern': r'\w+'},
- {'op': ragged.regex_replace,
- 'x': ragged.constant_value([['hello', '123'], ['1+1']]),
- 'pattern': r'\d',
- 'rewrite': '#'},
- {'op': ragged.substr,
- 'x': ragged.constant_value([['hello', '123'], ['1+1']]),
- 'pos': 2, 'len': 3},
- {'op': ragged.check_numerics,
- 'x': ragged.constant_value([[-2.0, 3.0], [-3.0]]),
- 'message': 'check-numerics'},
- ]
- ) # pyformat: disable
- def testUnaryOp(self, x, op=ragged.abs, **extra_args):
- x = ragged.convert_to_tensor_or_ragged_tensor(x)
- result = op(x, **extra_args)
-
- # Run the wrapped op on the dense values, for comparison.
- dense_x = x.inner_values if isinstance(x, ragged.RaggedTensor) else x
- expected_flat_values = array_ops.reshape(
- op.__wrapped__(dense_x, **extra_args), [-1])
-
- with self.test_session():
- # Check that the result has the expected shape.
- self.assertSameShape(x, result)
-
- # Check that the result has the expected (flattened) values.
- if isinstance(result, ragged.RaggedTensor):
- result_flat_values = array_ops.reshape(result.inner_values, [-1])
- else:
- result_flat_values = array_ops.reshape(result, [-1])
- self.assertAllEqual(expected_flat_values, result_flat_values)
-
- @parameterized.parameters(
- [
- #=====================================================================
- # Without broadcasting -- i.e., shapes match exactly.
- #=====================================================================
- # Shapes: x:(), y:()
- {'x': 12,
- 'y': 8},
- # Shapes: x:(3,), y:(3,)
- {'x': [7, 8, 9],
- 'y': [1, -2, 3]},
- # Shapes: x:(2, 2), y:(2, 2)
- {'x': [[-2, 3], [-3, -4]],
- 'y': [[1, 2], [3, 4]]},
- # Shapes: x:(2, None), y:(2, None)
- {'x': ragged.constant_value([[-2, 3], [-3]]),
- 'y': ragged.constant_value([[5, 6], [7]])},
- # Shapes: x:(2, 2, 2), y:(2, 2, 2)
- {'x': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
- 'y': [[[9, 3], [3, 4]], [[5, 2], [7, 6]]]},
- # Shapes: x:(2, None, None), y: (2, None, None)
- {'x': ragged.constant_value([[[1, 2], [3], [4]], [[], [5, 7, 8]]]),
- 'y': ragged.constant_value([[[3, 8], [2], [5]], [[], [1, 9, 8]]])},
- # Shapes: x:(2, None, 2), y: (2, None, 2)
- {'x': ragged.constant_value([[[1, 2]], [[3, 4], [5, 6], [7, 8]]],
- ragged_rank=1),
- 'y': ragged.constant_value([[[9, 3]], [[5, 2], [3, 4], [7, 6]]],
- ragged_rank=1)},
-
- #=====================================================================
- # With broadcasting
- #=====================================================================
- # Shapes: x:(), y:(3,)
- {'x': 12, # Broadcast () -> (3,)
- 'y': [1, -2, 3]},
- # Shapes: x:(1,), y:(3,)
- {'x': [12], # Broadcast (1,) -> (3,)
- 'y': [1, -2, 3]},
- # Shapes: x:(), y:(2, 2)
- {'x': 12, # Broadcast () -> (2, 2)
- 'y': [[1, 2], [3, 4]]},
- # Shapes: x:(1,), y:(2, 2)
- {'x': 12, # Broadcast (1,) -> (2, 2)
- 'y': [[1, 2], [3, 4]]},
- # Shapes: x:(2, 1), y:(2, 2)
- {'x': [[10], [20]], # Broadcast (2, 1) -> (2, 2)
- 'y': [[1, 2], [3, 4]]},
- # Shapes: x:(), y:(2, None)
- {'x': 10, # Broadcast () -> (2, None)
- 'y': ragged.constant_value([[1, 2], [3]], dtype=np.int32)},
- # TODO(edloper): Add tests for more advanced broadcasting, once we add
- # support for it.
-
- #=====================================================================
- # Keyword Args
- #=====================================================================
- {'x': ragged.constant_value([[[1, 2], [3], [4]], [[], [5, 7, 8]]]),
- 'y': ragged.constant_value([[[3, 8], [2], [5]], [[], [1, 9, 8]]]),
- 'use_kwargs': True},
- {'x': ragged.constant_value([[[1, 2]], [[3, 4], [5, 6], [7, 8]]],
- ragged_rank=1),
- 'y': ragged.constant_value([[[9, 3]], [[5, 2], [3, 4], [7, 6]]],
- ragged_rank=1),
- 'use_kwargs': True},
- ] +
- #=========================================================================
- # Test each unary op.
- #=========================================================================
- [{'x': ragged.constant_value([[-2.0, 3.0], [-3.0]]),
- 'y': ragged.constant_value([[5.0, 1.0], [12.0]]),
- 'op': op}
- for op in BINARY_FLOAT_OPS] +
- [{'x': ragged.constant_value([[-2, 3], [-3]]),
- 'y': ragged.constant_value([[5, 1], [12]]),
- 'op': op}
- for op in BINARY_INT_OPS] +
- [{'x': ragged.constant_value([[True, True], [False]]),
- 'y': ragged.constant_value([[False, True], [False]]),
- 'op': op}
- for op in BINARY_BOOL_OPS] +
- [
- ]
- ) # pyformat: disable
- def testBinaryOp(self, x, y, op=ragged.add, **extra_args):
- use_kwargs = extra_args.pop('use_kwargs', False)
- x = ragged.convert_to_tensor_or_ragged_tensor(x)
- y = ragged.convert_to_tensor_or_ragged_tensor(y)
- if use_kwargs:
- result = op(x=x, y=y, **extra_args)
- else:
- result = op(x, y, **extra_args)
-
- # Run the wrapped op on the dense values, for comparison.
- dense_x = x.inner_values if isinstance(x, ragged.RaggedTensor) else x
- dense_y = y.inner_values if isinstance(y, ragged.RaggedTensor) else y
- expected_flat_values = array_ops.reshape(
- op.__wrapped__(dense_x, dense_y, **extra_args), [-1])
-
- with self.test_session():
- # Check that the result has the expected shape.
- self.assertSameShape(y, result)
-
- # Check that the result has the expected (flattened) values.
- if isinstance(result, ragged.RaggedTensor):
- result_flat_values = array_ops.reshape(result.inner_values, [-1])
- else:
- result_flat_values = array_ops.reshape(result, [-1])
- self.assertAllEqual(expected_flat_values, result_flat_values)
-
- @parameterized.parameters(
- [
- {'inputs': (12, 8, 3)},
- {'inputs': ([1, 2, 3], [7, 8, 9], [3, 6, 9])},
- {'inputs': ([[1, 2]], [[3, 4]], [[5, 6]])},
- {'inputs': (ragged.constant_value([[1, 3], [-3]]),
- ragged.constant_value([[4, 7], [88]]),
- ragged.constant_value([[2, 9], [12]]))},
- {'inputs': (ragged.constant_value([[[1, 3], [-3]], [[1]]]),
- ragged.constant_value([[[4, 7], [88]], [[2]]]),
- ragged.constant_value([[[2, 9], [12]], [[8]]]))},
- {'inputs': (ragged.constant_value([[[1, 3], [3, 4]], [[1, 5]]],
- ragged_rank=1),
- ragged.constant_value([[[4, 7], [1, 2]], [[2, 2]]],
- ragged_rank=1),
- ragged.constant_value([[[2, 9], [5, 2]], [[8, 0]]],
- ragged_rank=1))},
- {'inputs': (ragged.constant_value([[[1, 3], [-3]], [[1]]]),
- ragged.constant_value([[[4, 7], [88]], [[2]]]),
- ragged.constant_value([[[2, 9], [12]], [[8]]])),
- 'use_kwargs': True},
- ] + [
- {'op': ragged.add_n,
- 'inputs': (ragged.constant_value([[1, 3], [-3]]),
- ragged.constant_value([[4, 7], [88]]),
- ragged.constant_value([[2, 9], [12]]))},
- {'op': ragged.string_join,
- 'inputs': (ragged.constant_value([['a', 'b'], ['c']]),
- ragged.constant_value([['foo', 'bar'], ['baz']]),
- ragged.constant_value([['2', '9'], ['12']]))},
- ]) # pyformat: disable
- def testListValuedOp(self, inputs, op=ragged.add_n, **extra_args):
- use_kwargs = extra_args.pop('use_kwargs', False)
- inputs = [ragged.convert_to_tensor_or_ragged_tensor(x) for x in inputs]
- if use_kwargs:
- result = op(inputs=inputs, **extra_args)
- else:
- result = op(inputs, **extra_args)
-
- # Run the wrapped op on the dense values, for comparison.
- dense_inputs = [
- x.inner_values if isinstance(x, ragged.RaggedTensor) else x
- for x in inputs
- ]
- expected_flat_values = array_ops.reshape(
- op.__wrapped__(dense_inputs, **extra_args), [-1])
-
- with self.test_session():
- # Check that the result has the expected shape.
- self.assertSameShape(inputs[0], result)
-
- # Check that the result has the expected (flattened) values.
- if isinstance(result, ragged.RaggedTensor):
- result_flat_values = array_ops.reshape(result.inner_values, [-1])
- else:
- result_flat_values = array_ops.reshape(result, [-1])
- self.assertAllEqual(expected_flat_values, result_flat_values)
-
- @test_util.run_deprecated_v1
- def testUnknownRankError(self):
- x = ragged.constant([[1, 2], [3]])
- y = ragged.from_row_splits(
- array_ops.placeholder_with_default([1, 2, 3], shape=None), x.row_splits)
- with self.assertRaisesRegexp(
- ValueError, r'Unable to broadcast: unknown rank'):
- ragged.add(x, y)
-
- @parameterized.parameters([
- dict(
- x=ragged.constant_value([[1, 2], [3]]),
- y=[[10]],
- expected=[[11, 12], [13]]),
- dict(
- x=ragged.constant_value([[[1, 2], [3, 4]], [[5]]], ragged_rank=2),
- y=ragged.constant_value([[[10], [20]], [[30]]], ragged_rank=1),
- expected=[[[11, 12], [23, 24]], [[35]]]),
- dict(
- x=ragged.constant_value([[[1]]]),
- y=ragged.constant_value([[1]]),
- expected=[[[2]]]),
- ])
- def testBroadcastAdd(self, x, y, expected):
- x = ragged.convert_to_tensor_or_ragged_tensor(x, dtype=dtypes.int32)
- y = ragged.convert_to_tensor_or_ragged_tensor(y, dtype=dtypes.int32)
- result = x + y
- with self.cached_session():
- self.assertEqual(result.eval().tolist(), expected)
-
- def testShapeMismatch(self):
- x = ragged.constant([[1, 2, 3], [4, 5]])
- y = ragged.constant([[1, 2, 3], [4, 5, 6]])
- with self.assertRaisesRegexp(errors.InvalidArgumentError,
- 'Incompatible shapes'):
- with self.cached_session():
- ragged.add(x, y).eval()
-
- def testDocstring(self):
- self.assertRegexpMatches(
- ragged.add.__doc__,
- 'Ragged version of the elementwise operation `tf.math.add`')
- self.assertEqual(ragged.add.__name__, 'add')
-
-
-if __name__ == '__main__':
- googletest.main()
diff --git a/tensorflow/python/ops/ragged/ragged_expand_dims_op_test.py b/tensorflow/python/ops/ragged/ragged_expand_dims_op_test.py
index 3ff6697..c747bb3 100644
--- a/tensorflow/python/ops/ragged/ragged_expand_dims_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_expand_dims_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.expand_dims."""
+"""Tests for ragged_array_ops.expand_dims."""
from __future__ import absolute_import
from __future__ import division
@@ -21,16 +21,19 @@
from absl.testing import parameterized
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedExpandDimsOpTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedExpandDimsOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
# An example 4-d ragged tensor with shape [3, (D2), (D3), 2], and the
# expected result calling for expand_dims on each axis. c.f. the table of
- # expected result shapes in the ragged.expand_dims docstring.
+ # expected result shapes in the ragged_array_ops.expand_dims docstring.
EXAMPLE4D = [[[[1, 1], [2, 2]], [[3, 3]]],
[],
[[], [[4, 4], [5, 5], [6, 6]]]] # pyformat: disable
@@ -105,21 +108,19 @@
expected=EXAMPLE4D_EXPAND_AXIS[4],
expected_shape=[3, None, None, 2, 1]),
]) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedExpandDims(self,
rt_input,
axis,
expected,
ragged_rank=None,
expected_shape=None):
- rt = ragged.constant(rt_input, ragged_rank=ragged_rank)
- expanded = ragged.expand_dims(rt, axis=axis)
+ rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
+ expanded = ragged_array_ops.expand_dims(rt, axis=axis)
self.assertEqual(expanded.shape.ndims, rt.shape.ndims + 1)
if expected_shape is not None:
self.assertEqual(expanded.shape.as_list(), expected_shape)
- with self.test_session():
- self.assertEqual(expanded.eval().tolist(), expected)
+ self.assertRaggedEqual(expanded, expected)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_factory_ops.py b/tensorflow/python/ops/ragged/ragged_factory_ops.py
index d1f301b..8cda987 100644
--- a/tensorflow/python/ops/ragged/ragged_factory_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_factory_ops.py
@@ -21,18 +21,16 @@
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.framework import tensor_util
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.util.tf_export import tf_export
#===============================================================================
# Op to construct a constant RaggedTensor from a nested Python list.
#===============================================================================
+@tf_export("ragged.constant")
def constant(pylist, dtype=None, ragged_rank=None, inner_shape=None, name=None):
"""Constructs a constant RaggedTensor from a nested Python list.
@@ -56,8 +54,8 @@
`pylist`.
ragged_rank: An integer specifying the ragged rank of the returned
`RaggedTensor`. Must be nonnegative and less than `K`. Defaults to
- `max(0, K - 1)` if `inner_shape` is not specified. Defaults to
- `max(0, K - 1 - len(inner_shape))` if `inner_shape` is specified.
+ `max(0, K - 1)` if `inner_shape` is not specified. Defaults to `max(0, K
+ - 1 - len(inner_shape))` if `inner_shape` is specified.
inner_shape: A tuple of integers specifying the shape for individual inner
values in the returned `RaggedTensor`. Defaults to `()` if `ragged_rank`
is not specified. If `ragged_rank` is specified, then a default is chosen
@@ -72,17 +70,19 @@
ValueError: If the scalar values in `pylist` have inconsistent nesting
depth; or if ragged_rank or inner_shape are incompatible with `pylist`.
"""
- with ops.name_scope(name, 'RaggedConstant'):
- return _constant_value(from_row_splits, constant_op.constant, pylist, dtype,
- ragged_rank, inner_shape)
+ with ops.name_scope(name, "RaggedConstant"):
+ return _constant_value(ragged_tensor.RaggedTensor.from_row_splits,
+ constant_op.constant, pylist, dtype, ragged_rank,
+ inner_shape)
+@tf_export(v1=["ragged.constant_value"])
def constant_value(pylist, dtype=None, ragged_rank=None, inner_shape=None):
"""Constructs a RaggedTensorValue from a nested Python list.
- > Warning: This function returns a `RaggedTensorValue`, not a `RaggedTensor`.
- > If you wish to construct a constant `RaggedTensor`, use
- > [`ragged.constant(...)`](constant.md) instead.
+ Warning: This function returns a `RaggedTensorValue`, not a `RaggedTensor`.
+ If you wish to construct a constant `RaggedTensor`, use
+ [`ragged.constant(...)`](constant.md) instead.
Example:
@@ -153,29 +153,29 @@
depth; or if ragged_rank or inner_shape are incompatible with `pylist`.
"""
if ragged_tensor.is_ragged(pylist):
- raise TypeError('pylist may not be a RaggedTensor or RaggedTensorValue.')
+ raise TypeError("pylist may not be a RaggedTensor or RaggedTensorValue.")
if not isinstance(pylist, (list, tuple)):
# Scalar value
if ragged_rank is not None and ragged_rank != 0:
- raise ValueError('Invalid pylist=%r: incompatible with ragged_rank=%d' %
+ raise ValueError("Invalid pylist=%r: incompatible with ragged_rank=%d" %
(pylist, ragged_rank))
if inner_shape is not None and inner_shape:
raise ValueError(
- 'Invalid pylist=%r: incompatible with dim(inner_shape)=%d' %
+ "Invalid pylist=%r: incompatible with dim(inner_shape)=%d" %
(pylist, len(inner_shape)))
return inner_factory(pylist, dtype, ())
if ragged_rank is not None and ragged_rank < 0:
raise ValueError(
- 'Invalid ragged_rank=%r: must be nonnegative' % ragged_rank)
+ "Invalid ragged_rank=%r: must be nonnegative" % ragged_rank)
# Find the depth of scalar values in `pylist`.
scalar_depth, max_depth = _find_scalar_and_max_depth(pylist)
if scalar_depth is not None:
if max_depth > scalar_depth:
- raise ValueError('Invalid pylist=%r: empty list nesting is greater '
- 'than scalar value nesting' % pylist)
+ raise ValueError("Invalid pylist=%r: empty list nesting is greater "
+ "than scalar value nesting" % pylist)
# If both inner_shape and ragged_rank were specified, then check that
# they are compatible with pylist.
@@ -184,8 +184,8 @@
if ((scalar_depth is not None and expected_depth != scalar_depth) or
(scalar_depth is None and expected_depth < max_depth)):
raise ValueError(
- 'Invalid pylist=%r: incompatible with ragged_rank=%d '
- 'and dim(inner_shape)=%d' % (pylist, ragged_rank, len(inner_shape)))
+ "Invalid pylist=%r: incompatible with ragged_rank=%d "
+ "and dim(inner_shape)=%d" % (pylist, ragged_rank, len(inner_shape)))
# Check if the result is a `Tensor`.
if (ragged_rank == 0 or
@@ -221,7 +221,7 @@
values = concatenated_values
values = inner_factory(
- values, dtype=dtype, shape=(len(values),) + inner_shape, name='values')
+ values, dtype=dtype, shape=(len(values),) + inner_shape, name="values")
for row_splits in reversed(nested_splits):
values = ragged_factory(values, row_splits)
return values
@@ -249,7 +249,7 @@
child_scalar_depth, child_max_depth = _find_scalar_and_max_depth(child)
if child_scalar_depth is not None:
if scalar_depth is not None and scalar_depth != child_scalar_depth + 1:
- raise ValueError('all scalar values must have the same nesting depth')
+ raise ValueError("all scalar values must have the same nesting depth")
scalar_depth = child_scalar_depth + 1
max_depth = max(max_depth, child_max_depth + 1)
return (scalar_depth, max_depth)
@@ -273,436 +273,24 @@
"""Checks that `item` has a consistent shape matching `shape`."""
is_nested = isinstance(item, (list, tuple))
if is_nested != bool(shape):
- raise ValueError('inner values have inconsistent shape')
+ raise ValueError("inner values have inconsistent shape")
if is_nested:
if shape[0] != len(item):
- raise ValueError('inner values have inconsistent shape')
+ raise ValueError("inner values have inconsistent shape")
for child in item:
check_inner_shape(child, shape[1:])
# Collapse the ragged layers to get the list of inner values.
- inner_values = pylist
+ flat_values = pylist
for dim in range(ragged_rank):
- if not all(isinstance(v, (list, tuple)) for v in inner_values):
- raise ValueError('pylist has scalar values depth %d, but ragged_rank=%d '
- 'requires scalar value depth greater than %d' %
+ if not all(isinstance(v, (list, tuple)) for v in flat_values):
+ raise ValueError("pylist has scalar values depth %d, but ragged_rank=%d "
+ "requires scalar value depth greater than %d" %
(dim + 1, ragged_rank, ragged_rank))
- inner_values = sum((list(v) for v in inner_values), [])
+ flat_values = sum((list(v) for v in flat_values), [])
# Compute the inner shape looking only at the leftmost elements; and then
# use check_inner_shape to verify that other elements have the same shape.
- inner_shape = get_inner_shape(inner_values)
- check_inner_shape(inner_values, inner_shape)
+ inner_shape = get_inner_shape(flat_values)
+ check_inner_shape(flat_values, inner_shape)
return inner_shape[1:]
-
-
-#===============================================================================
-# Convert value -> tensor
-#===============================================================================
-def convert_to_tensor_or_ragged_tensor(value,
- dtype=None,
- preferred_dtype=None,
- name=None):
- """Converts value to a `RaggedTensor` or `Tensor`.
-
- * If `value` is a `RaggedTensor`, then return it as-is.
- * If `value` is a `RaggedTensorValue`, return a corresponding constant
- `RaggedTensor`.
- * Otherwise, use `convert_to_tensor` to convert `value` to a `Tensor`.
-
- Args:
- value: A `RaggedTensor`, a `RaggedTensorValue`, or an object whose type has
- a registered `Tensor` conversion function.
- dtype: Optional element type for the returned tensor. If missing the type
- is inferred from the type of `value`.
- preferred_dtype: Optional element type for the returned tensor, used when
- dtype is None. This argument has no effect if `value` is already a
- tensor, or when conversion is not possible.
- name: Optional name to use if a new `Tensor` is created.
-
- Returns:
- A `Tensor` or `RaggedTensor`.
- """
- if isinstance(value, ragged_tensor.RaggedTensor):
- if dtype and not dtype.is_compatible_with(value.dtype):
- raise ValueError('Tensor conversion requested dtype %s for '
- 'RaggedTensor with dtype %s: %r' %
- (dtype.name, value.dtype.name, value))
- return value
- elif isinstance(value, ragged_tensor_value.RaggedTensorValue):
- with ops.name_scope(name, 'ConvertToTensorOrRaggedTensor', []):
- inner_values = ops.convert_to_tensor(
- value=value.inner_values,
- dtype=dtype,
- preferred_dtype=preferred_dtype,
- name='inner_values')
- return from_nested_row_splits(inner_values, value.nested_row_splits)
- else:
- return ops.convert_to_tensor(
- value=value, dtype=dtype, preferred_dtype=preferred_dtype, name=name)
-
-
-#===============================================================================
-# Ops to construct RaggedTensor from row-partitioned values.
-#===============================================================================
-
-
-def from_value_rowids(values, value_rowids, nrows=None, name=None):
- """Creates a `RaggedTensor` with rows partitioned by `value_rowids`.
-
- The returned `RaggedTensor` corresponds with the python list defined by:
-
- ```python
- result = [[values[i] for i in range(len(values)) if value_rowids[i] == row]
- for row in range(nrows)]
- ```
-
- Warning: currently, this needs to cast value_rowids to int64 before
- converting, since `tf.bincount` only supports `int32`.
-
- Args:
- values: A potentially ragged tensor with shape `[nvals, ...]`.
- value_rowids: A 1-D int64 tensor with shape `[nvals]`, which corresponds
- one-to-one with `values`, and specifies each value's row index. Must be
- nonnegative, and must be sorted in ascending order.
- nrows: An int64 scalar specifying the number of rows. This should be
- specified if the `RaggedTensor` may containing empty training rows. Must
- be greater than `value_rowids[-1]` (or zero if `value_rowids` is empty).
- Defaults to `value_rowids[-1]` (or zero if `value_rowids` is empty).
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor`. `result.rank = values.rank + 1`.
- `result.ragged_rank = values.ragged_rank + 1`.
-
- Raises:
- ValueError: If `nrows` is incompatible with `value_rowids`.
-
- #### Example:
- ```python
- >>> rt = ragged.from_value_rowids(
- ... values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... value_rowids=[0, 0, 0, 0, 2, 2, 2, 3],
- ... nrows=5)
- >>> rt.eval().tolist()
- [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
- ```
- """
- with ops.name_scope(name, 'RaggedFromValueRowIds',
- [values, value_rowids, nrows]):
- values = convert_to_tensor_or_ragged_tensor(values, name='values')
- value_rowids = ops.convert_to_tensor(
- value_rowids, dtypes.int64, name='value_rowids')
- if nrows is None:
- const_rowids = tensor_util.constant_value(value_rowids)
- if const_rowids is None:
- nrows = array_ops.concat([value_rowids[-1:], [-1]], axis=0)[0] + 1
- const_nrows = None
- else:
- const_nrows = const_rowids[-1] + 1 if const_rowids.size > 0 else 0
- nrows = ops.convert_to_tensor(const_nrows, dtypes.int64, name='nrows')
- else:
- nrows = ops.convert_to_tensor(nrows, dtypes.int64, 'nrows')
- const_nrows = tensor_util.constant_value(nrows)
- if const_nrows is not None:
- if const_nrows < 0:
- raise ValueError('Expected nrows >= 0; got %d' % const_nrows)
- const_rowids = tensor_util.constant_value(value_rowids)
- if const_rowids is not None and const_rowids.size > 0:
- if not const_nrows >= const_rowids[-1] + 1:
- raise ValueError(
- 'Expected nrows >= value_rowids[-1] + 1; got nrows=%d, '
- 'value_rowids[-1]=%d' % (const_nrows, const_rowids[-1]))
-
- value_rowids.shape.assert_has_rank(1)
- nrows.shape.assert_has_rank(0)
- values.shape[:1].assert_is_compatible_with(value_rowids.shape)
-
- # Convert value_rowids & nrows to row_splits.
- # Note: we don't use segment_ids_to_row_splits() here because we want
- # to save the intermediate value `row_lengths`, so we can cache it.
- # TODO(b/116708836) Upgrade bincount to accept int64 so we can skip the cast
- # (Remove the warning in the docstring when we do.)
- value_rowids_int32 = math_ops.cast(value_rowids, dtypes.int32)
- nrows_int32 = math_ops.cast(nrows, dtypes.int32)
- row_lengths = math_ops.bincount(
- value_rowids_int32,
- minlength=nrows_int32,
- maxlength=nrows_int32,
- dtype=dtypes.int64)
- row_splits = array_ops.concat([[0], math_ops.cumsum(row_lengths)], axis=0)
- if const_nrows is not None:
- row_lengths.set_shape([const_nrows])
- row_splits.set_shape([const_nrows + 1])
-
- return ragged_tensor.RaggedTensor(
- values,
- row_splits,
- cached_row_lengths=row_lengths,
- cached_value_rowids=value_rowids,
- cached_nrows=nrows,
- internal=True)
-
-
-def from_row_splits(values, row_splits, name=None):
- """Creates a `RaggedTensor` with rows partitioned by `row_splits`.
-
- The returned `RaggedTensor` corresponds with the python list defined by:
-
- ```python
- result = [values[row_splits[i]:row_splits[i + 1]]
- for i in range(len(row_splits) - 1)]
- ```
-
- Args:
- values: A potentially ragged tensor with shape `[nvals, ...]`.
- row_splits: A 1-D int64 tensor with shape `[nrows+1]`. Must not be empty,
- and must be sorted in ascending order. `row_splits[0]` must be zero and
- `row_splits[-1]` must be `nvals`.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor`. `result.rank = values.rank + 1`.
- `result.ragged_rank = values.ragged_rank + 1`.
-
- Raises:
- ValueError: If `row_splits` is an empty list.
-
- #### Example:
- ```python
- >>> rt = ragged.from_row_splits(
- ... values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... row_splits=[0, 4, 4, 7, 8, 8])
- >>> rt.eval().tolist()
- [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
- ```
- """
- if isinstance(row_splits, (list, tuple)) and not row_splits:
- raise ValueError('row_splits tensor may not be empty.')
- with ops.name_scope(name, 'RaggedFromRowSplits', [values, row_splits]):
- values = convert_to_tensor_or_ragged_tensor(values, name='values')
- row_splits = ops.convert_to_tensor(row_splits, dtypes.int64, 'row_splits')
- row_splits.shape.assert_has_rank(1)
- return ragged_tensor.RaggedTensor(
- values=values, row_splits=row_splits, internal=True)
-
-
-def from_row_lengths(values, row_lengths, name=None):
- """Creates a `RaggedTensor` with rows partitioned by `row_lengths`.
-
- The returned `RaggedTensor` corresponds with the python list defined by:
-
- ```python
- result = [[values.pop(0) for i in range(length)]
- for length in row_lengths]
- ```
-
- Args:
- values: A potentially ragged tensor with shape `[nvals, ...]`.
- row_lengths: A 1-D int64 tensor with shape `[nrows]`. Must be nonnegative.
- `sum(row_lengths)` must be `nvals`.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor`. `result.rank = values.rank + 1`.
- `result.ragged_rank = values.ragged_rank + 1`.
-
- #### Example:
- ```python
- >>> rt = ragged.from_row_lengths(
- ... values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... row_lengths=[4, 0, 3, 1, 0])
- >>> rt.eval().tolist()
- [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
- ```
- """
- with ops.name_scope(name, 'RaggedFromRowLengths', [values, row_lengths]):
- values = convert_to_tensor_or_ragged_tensor(values, name='values')
- row_lengths = ops.convert_to_tensor(row_lengths, dtypes.int64,
- 'row_lengths')
- row_lengths.shape.assert_has_rank(1)
- row_limits = math_ops.cumsum(row_lengths)
- row_splits = array_ops.concat([[0], row_limits], axis=0)
- return ragged_tensor.RaggedTensor(
- values=values,
- row_splits=row_splits,
- cached_row_lengths=row_lengths,
- internal=True)
-
-
-def from_row_starts(values, row_starts, name=None):
- """Creates a `RaggedTensor` with rows partitioned by `row_starts`.
-
- Equivalent to: `from_row_splits(values, concat([row_starts, nvals]))`.
-
- Args:
- values: A potentially ragged tensor with shape `[nvals, ...]`.
- row_starts: A 1-D int64 tensor with shape `[nrows]`. Must be nonnegative
- and sorted in ascending order. If `nrows>0`, then `row_starts[0]` must be
- zero.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor`. `result.rank = values.rank + 1`.
- `result.ragged_rank = values.ragged_rank + 1`.
-
- #### Example:
- ```python
- >>> rt = ragged.from_row_starts(
- ... values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... row_starts=[0, 4, 4, 7, 8])
- >>> rt.eval().tolist()
- [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
- ```
- """
- with ops.name_scope(name, 'RaggedFromRowStarts', [values, row_starts]):
- values = convert_to_tensor_or_ragged_tensor(values, name='values')
- row_starts = ops.convert_to_tensor(row_starts, dtypes.int64, 'row_starts')
- row_starts.shape.assert_has_rank(1)
- nvals = array_ops.shape(values, out_type=dtypes.int64)[:1]
- row_splits = array_ops.concat([row_starts, nvals], axis=0)
- return ragged_tensor.RaggedTensor(
- values=values, row_splits=row_splits, internal=True)
-
-
-def from_row_limits(values, row_limits, name=None):
- """Creates a `RaggedTensor` with rows partitioned by `row_limits`.
-
- Equivalent to: `from_row_splits(values, concat([0, row_limits]))`.
-
- Args:
- values: A potentially ragged tensor with shape `[nvals, ...]`.
- row_limits: A 1-D int64 tensor with shape `[nrows]`. Must be sorted in
- ascending order. If `nrows>0`, then `row_limits[-1]` must be `nvals`.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor`. `result.rank = values.rank + 1`.
- `result.ragged_rank = values.ragged_rank + 1`.
-
- #### Example:
- ```python
- >>> rt = ragged.from_row_limits(
- ... values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... row_limits=[4, 4, 7, 8, 8])
- >>> rt.eval().tolist()
- [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
- ```
- """
- with ops.name_scope(name, 'RaggedFromRowLimits', [values, row_limits]):
- values = convert_to_tensor_or_ragged_tensor(values, name='values')
- row_limits = ops.convert_to_tensor(row_limits, dtypes.int64, 'row_limits')
- row_limits.shape.assert_has_rank(1)
- zero = array_ops.zeros([1], dtypes.int64)
- row_splits = array_ops.concat([zero, row_limits], axis=0)
- return ragged_tensor.RaggedTensor(
- values=values, row_splits=row_splits, internal=True)
-
-
-def from_nested_value_rowids(inner_values,
- nested_value_rowids,
- nested_nrows=None,
- name=None):
- """Creates a `RaggedTensor` from a nested list of `value_rowids` tensors.
-
- Equivalent to:
-
- ```python
- result = inner_values
- for (value_rowids, nrows) in reversed(zip(nested_value_rowids, nested_nrows)):
- result = from_value_rowids(result, value_rowids, nrows)
- ```
-
- Args:
- inner_values: A potentially ragged tensor.
- nested_value_rowids: A list of 1-D int64 tensors. The `i`th tensor is used
- as the `value_rowids` for the `i`th ragged dimension.
- nested_nrows: A list of int64 scalars. The `i`th scalar is used as the
- `nrows` for the `i`th ragged dimension.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor` (or `inner_values` if `nested_value_rowids` is empty).
-
- Raises:
- ValueError: If `len(nested_values_rowids) != len(nested_nrows)`.
- """
- if isinstance(nested_value_rowids, ops.Tensor):
- raise TypeError('nested_value_rowids must be a list of Tensors')
- if nested_nrows is None:
- nested_nrows = [None] * len(nested_value_rowids)
- else:
- if isinstance(nested_nrows, ops.Tensor):
- raise TypeError('nested_nrows must be a list of Tensors')
- if len(nested_nrows) != len(nested_value_rowids):
- raise ValueError('nested_nrows must have the same length as '
- 'nested_value_rowids')
-
- with ops.name_scope(
- name, 'RaggedFromNestedValueRowIds',
- [inner_values] + list(nested_value_rowids) + list(nested_nrows)):
- result = inner_values
- for value_rowids, nrows in reversed(
- list(zip(nested_value_rowids, nested_nrows))):
- result = from_value_rowids(result, value_rowids, nrows)
- return result
-
-
-def from_nested_row_splits(inner_values, nested_row_splits, name=None):
- """Creates a `RaggedTensor` from a nested list of `row_splits` tensors.
-
- Equivalent to:
-
- ```python
- result = inner_values
- for row_splits in reversed(nested_row_splits):
- result = from_row_splits(result, row_splits)
- ```
-
- Args:
- inner_values: A potentially ragged tensor.
- nested_row_splits: A list of 1-D int64 tensors. The `i`th tensor is used as
- the `row_splits` for the `i`th ragged dimension.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor` (or `inner_values` if `nested_row_splits` is empty).
- """
- if isinstance(nested_row_splits, ops.Tensor):
- raise TypeError('nested_row_splits must be a list of Tensors')
- with ops.name_scope(name, 'RaggedFromNestedRowSplits',
- [inner_values] + list(nested_row_splits)):
- result = inner_values
- for splits in reversed(nested_row_splits):
- result = from_row_splits(result, splits)
- return result
-
-
-def from_nested_row_lengths(inner_values, nested_row_lengths, name=None):
- """Creates a `RaggedTensor` from a nested list of `row_lengths` tensors.
-
- Equivalent to:
-
- ```python
- result = inner_values
- for row_lengths in reversed(nested_row_lengths):
- result = from_row_lengths(result, row_lengths)
- ```
-
- Args:
- inner_values: A potentially ragged tensor.
- nested_row_lengths: A list of 1-D int64 tensors. The `i`th tensor is used
- as the `row_lengths` for the `i`th ragged dimension.
- name: A name prefix for the RaggedTensor (optional).
-
- Returns:
- A `RaggedTensor` (or `inner_values` if `nested_row_lengths` is empty).
- """
- if isinstance(nested_row_lengths, ops.Tensor):
- raise TypeError('nested_row_lengths must be a list of Tensors')
- with ops.name_scope(name, 'RaggedFromNestedRowlengths',
- [inner_values] + list(nested_row_lengths)):
- result = inner_values
- for lengths in reversed(nested_row_lengths):
- result = from_row_lengths(result, lengths)
- return result
diff --git a/tensorflow/python/ops/ragged/ragged_from_sparse_op_test.py b/tensorflow/python/ops/ragged/ragged_from_sparse_op_test.py
index 3c0db9e..c6998e2 100644
--- a/tensorflow/python/ops/ragged/ragged_from_sparse_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_from_sparse_op_test.py
@@ -12,77 +12,77 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.from_sparse."""
+"""Tests for RaggedTensor.from_sparse."""
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 dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
from tensorflow.python.platform import googletest
-class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToSparseOpTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testDocStringExample(self):
st = sparse_tensor.SparseTensor(
indices=[[0, 0], [0, 1], [0, 2], [1, 0], [3, 0]],
values=[1, 2, 3, 4, 5],
dense_shape=[4, 3])
- rt = ragged.from_sparse(st)
+ rt = RaggedTensor.from_sparse(st)
- with self.test_session():
- self.assertEqual(rt.eval().tolist(), [[1, 2, 3], [4], [], [5]])
+ self.assertRaggedEqual(rt, [[1, 2, 3], [4], [], [5]])
- @test_util.run_deprecated_v1
def testEmpty(self):
st = sparse_tensor.SparseTensor(
indices=array_ops.zeros([0, 2], dtype=dtypes.int64),
values=[],
dense_shape=[4, 3])
- rt = ragged.from_sparse(st)
+ rt = RaggedTensor.from_sparse(st)
- with self.test_session():
- self.assertEqual(rt.eval().tolist(), [[], [], [], []])
+ self.assertRaggedEqual(rt, [[], [], [], []])
- @test_util.run_deprecated_v1
def testBadSparseTensorRank(self):
st1 = sparse_tensor.SparseTensor(indices=[[0]], values=[0], dense_shape=[3])
+ self.assertRaisesRegexp(ValueError, r'rank\(st_input\) must be 2',
+ RaggedTensor.from_sparse, st1)
+
st2 = sparse_tensor.SparseTensor(
indices=[[0, 0, 0]], values=[0], dense_shape=[3, 3, 3])
- st3 = sparse_tensor.SparseTensor(
- indices=array_ops.placeholder(dtypes.int64),
- values=[0],
- dense_shape=array_ops.placeholder(dtypes.int64))
self.assertRaisesRegexp(ValueError, r'rank\(st_input\) must be 2',
- ragged.from_sparse, st1)
- self.assertRaisesRegexp(ValueError, r'rank\(st_input\) must be 2',
- ragged.from_sparse, st2)
- self.assertRaisesRegexp(ValueError, r'rank\(st_input\) must be 2',
- ragged.from_sparse, st3)
+ RaggedTensor.from_sparse, st2)
- @test_util.run_deprecated_v1
+ if not context.executing_eagerly():
+ st3 = sparse_tensor.SparseTensor(
+ indices=array_ops.placeholder(dtypes.int64),
+ values=[0],
+ dense_shape=array_ops.placeholder(dtypes.int64))
+ self.assertRaisesRegexp(ValueError, r'rank\(st_input\) must be 2',
+ RaggedTensor.from_sparse, st3)
+
def testGoodPartialSparseTensorRank(self):
- st1 = sparse_tensor.SparseTensor(
- indices=[[0, 0]],
- values=[0],
- dense_shape=array_ops.placeholder(dtypes.int64))
- st2 = sparse_tensor.SparseTensor(
- indices=array_ops.placeholder(dtypes.int64),
- values=[0],
- dense_shape=[4, 3])
+ if not context.executing_eagerly():
+ st1 = sparse_tensor.SparseTensor(
+ indices=[[0, 0]],
+ values=[0],
+ dense_shape=array_ops.placeholder(dtypes.int64))
+ st2 = sparse_tensor.SparseTensor(
+ indices=array_ops.placeholder(dtypes.int64),
+ values=[0],
+ dense_shape=[4, 3])
- # Shouldn't throw ValueError
- ragged.from_sparse(st1)
- ragged.from_sparse(st2)
+ # Shouldn't throw ValueError
+ RaggedTensor.from_sparse(st1)
+ RaggedTensor.from_sparse(st2)
- @test_util.run_deprecated_v1
def testNonRaggedSparseTensor(self):
# "index_suffix" means the value of the innermost dimension of the index
# (i.e., indices[i][-1]).
@@ -92,22 +92,21 @@
# index_suffix of first index is not zero.
st1 = sparse_tensor.SparseTensor(
indices=[[0, 1], [0, 2], [2, 0]], values=[1, 2, 3], dense_shape=[3, 3])
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'.*SparseTensor is not right-ragged'):
+ self.evaluate(RaggedTensor.from_sparse(st1))
# index_suffix of an index that starts a new row is not zero.
st2 = sparse_tensor.SparseTensor(
indices=[[0, 0], [0, 1], [2, 1]], values=[1, 2, 3], dense_shape=[3, 3])
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'.*SparseTensor is not right-ragged'):
+ self.evaluate(RaggedTensor.from_sparse(st2))
# index_suffix of an index that continues a row skips a cell.
st3 = sparse_tensor.SparseTensor(
indices=[[0, 1], [0, 1], [0, 3]], values=[1, 2, 3], dense_shape=[3, 3])
- rt1 = ragged.from_sparse(st1)
- rt2 = ragged.from_sparse(st2)
- rt3 = ragged.from_sparse(st3)
- with self.test_session():
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'.*SparseTensor is not right-ragged', rt1.eval)
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'.*SparseTensor is not right-ragged', rt2.eval)
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'.*SparseTensor is not right-ragged', rt3.eval)
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'.*SparseTensor is not right-ragged'):
+ self.evaluate(RaggedTensor.from_sparse(st3))
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_from_tensor_op_test.py b/tensorflow/python/ops/ragged/ragged_from_tensor_op_test.py
index 1d8a00c..68d3953 100644
--- a/tensorflow/python/ops/ragged/ragged_from_tensor_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_from_tensor_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.from_tensor."""
+"""Tests for RaggedTensor.from_tensor."""
from __future__ import absolute_import
from __future__ import division
@@ -24,29 +24,26 @@
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
from tensorflow.python.platform import googletest
-class RaggedFromTensorOpTest(test_util.TensorFlowTestCase,
- parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToSparseOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
- @test_util.run_deprecated_v1
def testDocStringExamples(self):
- # The examples from ragged.from_tensor.__doc__.
+ # The examples from RaggedTensor.from_tensor.__doc__.
dt = constant_op.constant([[5, 7, 0], [0, 3, 0], [6, 0, 0]])
- with self.test_session():
- self.assertEqual(
- ragged.from_tensor(dt).eval().tolist(),
- [[5, 7, 0], [0, 3, 0], [6, 0, 0]])
+ self.assertRaggedEqual(
+ RaggedTensor.from_tensor(dt), [[5, 7, 0], [0, 3, 0], [6, 0, 0]])
- self.assertEqual(
- ragged.from_tensor(dt, lengths=[1, 0, 3]).eval().tolist(),
- [[5], [], [6, 0, 0]])
+ self.assertRaggedEqual(
+ RaggedTensor.from_tensor(dt, lengths=[1, 0, 3]), [[5], [], [6, 0, 0]])
- self.assertEqual(
- ragged.from_tensor(dt, padding=0).eval().tolist(),
- [[5, 7], [0, 3], [6]])
+ self.assertRaggedEqual(
+ RaggedTensor.from_tensor(dt, padding=0), [[5, 7], [0, 3], [6]])
@parameterized.parameters(
# 2D test cases, no length or padding.
@@ -263,7 +260,6 @@
[[[5, 6], [7]], [[0, 8], []]]]
},
) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedFromTensor(self,
tensor,
expected,
@@ -271,30 +267,27 @@
padding=None,
ragged_rank=1):
dt = constant_op.constant(tensor)
- rt = ragged.from_tensor(dt, lengths, padding, ragged_rank)
- self.assertEqual(type(rt), ragged.RaggedTensor)
+ rt = RaggedTensor.from_tensor(dt, lengths, padding, ragged_rank)
+ self.assertEqual(type(rt), RaggedTensor)
self.assertEqual(rt.ragged_rank, ragged_rank)
self.assertTrue(
dt.shape.is_compatible_with(rt.shape),
'%s is incompatible with %s' % (dt.shape, rt.shape))
- with self.test_session():
- self.assertEqual(rt.eval().tolist(), expected)
+ self.assertRaggedEqual(rt, expected)
- @test_util.run_deprecated_v1
def testHighDimensions(self):
# Use distinct prime numbers for all dimension shapes in this test, so
# we can see any errors that are caused by mixing up dimension sizes.
dt = array_ops.reshape(
math_ops.range(3 * 5 * 7 * 11 * 13 * 17), [3, 5, 7, 11, 13, 17])
for ragged_rank in range(1, 4):
- rt = ragged.from_tensor(dt, ragged_rank=ragged_rank)
- self.assertEqual(type(rt), ragged.RaggedTensor)
+ rt = RaggedTensor.from_tensor(dt, ragged_rank=ragged_rank)
+ self.assertEqual(type(rt), RaggedTensor)
self.assertEqual(rt.ragged_rank, ragged_rank)
self.assertTrue(
dt.shape.is_compatible_with(rt.shape),
'%s is incompatible with %s' % (dt.shape, rt.shape))
- with self.test_session():
- self.assertEqual(rt.eval().tolist(), self.evaluate(dt).tolist())
+ self.assertRaggedEqual(rt, self.evaluate(dt).tolist())
@parameterized.parameters(
# With no padding or lengths
@@ -398,15 +391,13 @@
'expected': [[], []]
},
)
- @test_util.run_deprecated_v1
def testEmpty(self, dt_shape, expected, lengths=None, padding=None):
dt = array_ops.zeros(dt_shape)
- rt = ragged.from_tensor(dt, lengths, padding)
- self.assertEqual(type(rt), ragged.RaggedTensor)
+ rt = RaggedTensor.from_tensor(dt, lengths, padding)
+ self.assertEqual(type(rt), RaggedTensor)
self.assertEqual(rt.ragged_rank, 1)
self.assertTrue(dt.shape.is_compatible_with(rt.shape))
- with self.test_session():
- self.assertEqual(rt.eval().tolist(), expected)
+ self.assertRaggedEqual(rt, expected)
@parameterized.parameters(
{
@@ -423,7 +414,7 @@
{
'tensor': [[1]],
'padding': 'a',
- 'error': (TypeError, "Expected int32, got 'a'.*")
+ 'error': (TypeError, '.*')
},
{
'tensor': [[1]],
@@ -451,7 +442,6 @@
'error': (ValueError, r'ragged_rank must be greater than 0; got -1')
},
)
- @test_util.run_deprecated_v1
def testErrors(self,
tensor,
lengths=None,
@@ -459,8 +449,8 @@
ragged_rank=1,
error=None):
dt = constant_op.constant(tensor)
- self.assertRaisesRegexp(error[0], error[1], ragged.from_tensor, dt, lengths,
- padding, ragged_rank)
+ self.assertRaisesRegexp(error[0], error[1], RaggedTensor.from_tensor, dt,
+ lengths, padding, ragged_rank)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_functional_ops.py b/tensorflow/python/ops/ragged/ragged_functional_ops.py
index 6b71d88..b6937a1 100644
--- a/tensorflow/python/ops/ragged/ragged_functional_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_functional_ops.py
@@ -19,40 +19,41 @@
from __future__ import print_function
from tensorflow.python.framework import ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.util.tf_export import tf_export
-def map_inner_values(op, *args, **kwargs):
- """Applies `op` to the inner values of one or more RaggedTensors.
+@tf_export("ragged.map_flat_values")
+def map_flat_values(op, *args, **kwargs):
+ """Applies `op` to the values of one or more RaggedTensors.
- Replaces any `RaggedTensor` in `args` or `kwargs` with its `inner_values`
+ Replaces any `RaggedTensor` in `args` or `kwargs` with its `flat_values`
tensor, and then calls `op`. Returns a `RaggedTensor` that is constructed
- from the input `RaggedTensor`s' `splits` and the value returned by
+ from the input `RaggedTensor`s' `nested_row_splits` and the value returned by
the `op`.
If the input arguments contain multiple `RaggedTensor`s, then they must have
- identical `splits`.
+ identical `nested_row_splits`.
Examples:
```python
>>> rt = ragged.constant([[1, 2, 3], [], [4, 5], [6]])
- >>> ragged.map_inner_values(tf.ones_like, rt).eval().tolist()
+ >>> ragged.map_flat_values(tf.ones_like, rt).eval().tolist()
[[1, 1, 1], [], [1, 1], [1]]
- >>> ragged.map_inner_values(tf.multiply, rt, rt).eval().tolist()
+ >>> ragged.map_flat_values(tf.multiply, rt, rt).eval().tolist()
[[1, 4, 9], [], [16, 25], [36]]
- >>> ragged.map_inner_values(tf.add, rt, 5).eval().tolist()
+ >>> ragged.map_flat_values(tf.add, rt, 5).eval().tolist()
[[6, 7, 8], [], [9, 10], [11]]
```
Args:
- op: The operation that should be applied to the RaggedTensor `inner_values`.
+ op: The operation that should be applied to the RaggedTensor `flat_values`.
`op` is typically an element-wise operation (such as math_ops.add), but
any operation that preserves the size of the outermost dimension can be
used. I.e., `shape[0]` of the value returned by `op` must match
- `shape[0]` of the `RaggedTensor`s' `inner_values` tensors.
+ `shape[0]` of the `RaggedTensor`s' `flat_values` tensors.
*args: Arguments for `op`.
**kwargs: Keyword arguments for `op`.
@@ -66,8 +67,8 @@
# Replace RaggedTensors with their values; and collect the splits tensors
# from each RaggedTensor.
nested_splits_lists = []
- inner_args = _replace_ragged_with_inner_values(args, nested_splits_lists)
- inner_kwargs = _replace_ragged_with_inner_values(kwargs, nested_splits_lists)
+ inner_args = _replace_ragged_with_flat_values(args, nested_splits_lists)
+ inner_kwargs = _replace_ragged_with_flat_values(kwargs, nested_splits_lists)
if not nested_splits_lists:
return op(*args, **kwargs)
@@ -75,15 +76,15 @@
ragged_util.assert_splits_match(nested_splits_lists)):
# Delegate to op, and then compose the result from the transformed values
# and the splits.
- return ragged_factory_ops.from_nested_row_splits(
+ return ragged_tensor.RaggedTensor.from_nested_row_splits(
op(*inner_args, **inner_kwargs), nested_splits_lists[0])
-def _replace_ragged_with_inner_values(value, nested_splits_lists):
- """Replace RaggedTensors with their inner_values, and record their splits.
+def _replace_ragged_with_flat_values(value, nested_splits_lists):
+ """Replace RaggedTensors with their flat_values, and record their splits.
Returns a copy of `value`, with any nested `RaggedTensor`s replaced by their
- `inner_values` tensor. Looks inside lists, tuples, and dicts.
+ `flat_values` tensor. Looks inside lists, tuples, and dicts.
Appends each `RaggedTensor`'s `nested_splits` to `nested_splits_lists`.
@@ -97,13 +98,13 @@
"""
# Base case
if ragged_tensor.is_ragged(value):
- value = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(value)
+ value = ragged_tensor.convert_to_tensor_or_ragged_tensor(value)
nested_splits_lists.append(value.nested_row_splits)
- return value.inner_values
+ return value.flat_values
# Recursion cases
def recurse(v):
- return _replace_ragged_with_inner_values(v, nested_splits_lists)
+ return _replace_ragged_with_flat_values(v, nested_splits_lists)
if isinstance(value, list):
return [recurse(v) for v in value]
diff --git a/tensorflow/python/ops/ragged/ragged_gather_nd_op_test.py b/tensorflow/python/ops/ragged/ragged_gather_nd_op_test.py
index 62c6819..d4bffeb 100644
--- a/tensorflow/python/ops/ragged/ragged_gather_nd_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_gather_nd_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for tf.ragged.gather_nd."""
+"""Tests for ragged_array_ops.gather_nd."""
from __future__ import absolute_import
from __future__ import division
@@ -21,14 +21,19 @@
from absl.testing import parameterized
import numpy as np
+from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedGatherNdOpTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedGatherNdOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
DOCSTRING_PARAMS = [[['000', '001'], ['010']],
@@ -41,18 +46,19 @@
#=========================================================================
dict(
descr='Docstring example 1',
- params=ragged.constant_value(DOCSTRING_PARAMS),
+ params=ragged_factory_ops.constant_value(DOCSTRING_PARAMS),
indices=[[2], [0]],
- expected=ragged.constant_value([[[], [b'210']],
- [[b'000', b'001'], [b'010']]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[], [b'210']], [[b'000', b'001'], [b'010']]])),
dict(
descr='Docstring example 2',
- params=ragged.constant_value(DOCSTRING_PARAMS),
+ params=ragged_factory_ops.constant_value(DOCSTRING_PARAMS),
indices=[[2, 1], [0, 0]],
- expected=ragged.constant_value([[b'210'], [b'000', b'001']])),
+ expected=ragged_factory_ops.constant_value(
+ [[b'210'], [b'000', b'001']])),
dict(
descr='Docstring example 3',
- params=ragged.constant_value(DOCSTRING_PARAMS),
+ params=ragged_factory_ops.constant_value(DOCSTRING_PARAMS),
indices=[[0, 0, 1], [1, 1, 2]],
expected=[b'001', b'112']),
#=========================================================================
@@ -60,175 +66,179 @@
#=========================================================================
dict(
descr='params: [B1, (B2)], indices: [0], result: [B1, (B2)]',
- params=ragged.constant_value([['a', 'b', 'c'], ['d']]),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
indices=np.zeros([0], dtype=np.int32),
- expected=ragged.constant_value([[b'a', b'b', b'c'], [b'd']])),
+ expected=ragged_factory_ops.constant_value(
+ [[b'a', b'b', b'c'], [b'd']])),
dict(
descr='params: [B1, (B2)], indices: [A1, 0], result: [A1, B1, (B2)]',
- params=ragged.constant_value([['a', 'b', 'c'], ['d']]),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
indices=np.zeros([3, 0], dtype=np.int32),
- expected=ragged.constant_value([[[b'a', b'b', b'c'], [b'd']],
- [[b'a', b'b', b'c'], [b'd']],
- [[b'a', b'b', b'c'], [b'd']]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[b'a', b'b', b'c'], [b'd']],
+ [[b'a', b'b', b'c'], [b'd']],
+ [[b'a', b'b', b'c'], [b'd']]])),
dict(
descr=('params: [B1, (B2)], indices: [A1, A2, 0], '
'result: [A1, A2, B1, (B2)]'),
- params=ragged.constant_value([['a', 'b', 'c'], ['d']]),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
indices=np.zeros([1, 3, 0], dtype=np.int32),
- expected=ragged.constant_value([[[[b'a', b'b', b'c'], [b'd']],
- [[b'a', b'b', b'c'], [b'd']],
- [[b'a', b'b', b'c'], [b'd']]]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[[b'a', b'b', b'c'], [b'd']],
+ [[b'a', b'b', b'c'], [b'd']],
+ [[b'a', b'b', b'c'], [b'd']]]])),
dict(
descr='params: [B1], indices: [A1, (A2), 0], result: [A1, (A2), B1]',
params=['a'],
- indices=ragged.constant_value([[[], []], [[]]],
- ragged_rank=1,
- dtype=np.int32),
- expected=ragged.constant_value([[[b'a'], [b'a']], [[b'a']]],
- ragged_rank=1)),
+ indices=ragged_factory_ops.constant_value(
+ [[[], []], [[]]],
+ ragged_rank=1,
+ dtype=np.int32),
+ expected=ragged_factory_ops.constant_value(
+ [[[b'a'], [b'a']], [[b'a']]],
+ ragged_rank=1)),
#=========================================================================
# Indices with 1 value (selects row from params)
#=========================================================================
dict(
descr='params: [B1, (B2)], indices: [A1, 1], result: [A1, (B2)]',
- params=ragged.constant_value([['a', 'b', 'c'], ['d']]),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
indices=[[1], [0]],
- expected=ragged.constant_value([[b'd'], [b'a', b'b', b'c']])),
+ expected=ragged_factory_ops.constant_value(
+ [[b'd'], [b'a', b'b', b'c']])),
dict(
descr=('params: [B1, (B2), (B3)], indices: [A1, 1], '
'result: [A1, (B2), (B3)]'),
- params=ragged.constant_value([[['a', 'b', 'c'], ['d']],
- [['e', 'f']]]),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
indices=[[1], [1]],
- expected=ragged.constant_value([[[b'e', b'f']], [[b'e', b'f']]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[b'e', b'f']], [[b'e', b'f']]])),
dict(
descr=('params: [B1, B2, B3], indices: [A1, (A2), 1], '
'result: [A1, (A2), B2, B3]'),
params=[[['a']], [['b']]],
- indices=ragged.constant_value([[[0]]], ragged_rank=1),
- expected=ragged.constant_value([[[[b'a']]]], ragged_rank=1)),
+ indices=ragged_factory_ops.constant_value([[[0]]], ragged_rank=1),
+ expected=ragged_factory_ops.constant_value(
+ [[[[b'a']]]], ragged_rank=1)),
#=========================================================================
# Indices with 2 values (selects row & col from params)
#=========================================================================
dict(
descr='params: [B1, (B2)], indices: [A1, 2], result: [A1]',
- params=ragged.constant_value([['a', 'b', 'c'], ['d']]),
+ params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
indices=[[1, 0], [0, 0], [0, 2]],
- expected=ragged.constant_value([b'd', b'a', b'c'])),
+ expected=ragged_factory_ops.constant_value([b'd', b'a', b'c'])),
dict(
descr=('params: [B1, (B2), (B3)], indices: [A1, 2], '
'result: [A1, (B3)]'),
- params=ragged.constant_value([[['a', 'b', 'c'], ['d']],
- [['e', 'f']]]),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
indices=[[1, 0], [0, 1], [0, 0]],
- expected=ragged.constant_value([[b'e', b'f'], [b'd'],
- [b'a', b'b', b'c']])),
+ expected=ragged_factory_ops.constant_value(
+ [[b'e', b'f'], [b'd'], [b'a', b'b', b'c']])),
dict(
descr=('params: [B1, (B2), (B3)], indices: [A1, A2, 2], '
'result: [A1, (A2), (B3)]'),
- params=ragged.constant_value([[['a', 'b', 'c'], ['d']],
- [['e', 'f']]]),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
indices=[[[1, 0], [0, 1], [0, 0]]],
- expected=ragged.constant_value([[[b'e', b'f'], [b'd'],
- [b'a', b'b', b'c']]])),
+ expected=ragged_factory_ops.constant_value(
+ [[[b'e', b'f'], [b'd'], [b'a', b'b', b'c']]])),
dict(
descr=('params: [B1, (B2), B3], indices: [A1, A2, 2], '
'result: [A1, A2, B3]'),
- params=ragged.constant_value([[['a', 'b'], ['c', 'd']],
- [['e', 'f']]],
- ragged_rank=1),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b'], ['c', 'd']],
+ [['e', 'f']]],
+ ragged_rank=1),
indices=[[[1, 0], [0, 1], [0, 0]]],
expected=[[[b'e', b'f'], [b'c', b'd'], [b'a', b'b']]]),
dict(
descr=('params: [B1, (B2), B3], indices: [A1, A2, A3, 2], '
'result: [A1, A2, A3, B3]'),
- params=ragged.constant_value([[['a', 'b'], ['c', 'd']],
- [['e', 'f']]],
- ragged_rank=1),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b'], ['c', 'd']],
+ [['e', 'f']]],
+ ragged_rank=1),
indices=[[[[1, 0], [0, 1], [0, 0]]]],
expected=[[[[b'e', b'f'], [b'c', b'd'], [b'a', b'b']]]]),
dict(
descr=('params: [B1, (B2), (B3)], indices: [A1, (A2), 2], '
'result: [A1, (A2), (B3)]'),
- params=ragged.constant_value([[['a', 'b', 'c'], ['d']],
- [['e', 'f']]]),
- indices=ragged.constant_value([[[1, 0], [0, 1]], [[0, 0]]],
- ragged_rank=1),
- expected=ragged.constant_value([[[b'e', b'f'], [b'd']],
- [[b'a', b'b', b'c']]])),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
+ indices=ragged_factory_ops.constant_value(
+ [[[1, 0], [0, 1]], [[0, 0]]],
+ ragged_rank=1),
+ expected=ragged_factory_ops.constant_value(
+ [[[b'e', b'f'], [b'd']], [[b'a', b'b', b'c']]])),
#=========================================================================
# Indices with 3 values
#=========================================================================
dict(
descr=('params: [B1, (B2), (B3)], indices: [A1, 3], '
'result: [A1]'),
- params=ragged.constant_value([[['a', 'b', 'c'], ['d']],
- [['e', 'f']]]),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
indices=[[1, 0, 1], [0, 0, 0], [0, 1, 0]],
expected=[b'f', b'a', b'd']),
dict(
descr=('params: [B1, (B2), B3], indices: [A1, 3], '
'result: [A1]'),
- params=ragged.constant_value([[['a', 'b'], ['c', 'd']],
- [['e', 'f']]],
- ragged_rank=1),
+ params=ragged_factory_ops.constant_value(
+ [[['a', 'b'], ['c', 'd']], [['e', 'f']]],
+ ragged_rank=1),
indices=[[1, 0, 1], [0, 0, 0], [0, 1, 1]],
expected=[b'f', b'a', b'd']),
dict(
descr=('params: [B1, (B2), (B3), B4], indices: [A1, 3], '
'result: [A1, B4]'),
- params=ragged.constant_value([[[['a', 'b'], ['c', 'd']],
- [['e', 'f']]]],
- ragged_rank=2),
+ params=ragged_factory_ops.constant_value(
+ [[[['a', 'b'], ['c', 'd']], [['e', 'f']]]],
+ ragged_rank=2),
indices=[[0, 0, 1], [0, 0, 0], [0, 1, 0]],
expected=[[b'c', b'd'], [b'a', b'b'], [b'e', b'f']]),
]) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedGatherNd(self, descr, params, indices, expected):
- result = ragged.gather_nd(params, indices)
- self.assertEqual(
- getattr(result, 'ragged_rank', 0), getattr(expected, 'ragged_rank', 0))
- with self.test_session() as sess:
- if hasattr(expected, 'tolist'):
- expected = expected.tolist()
- self.assertEqual(self.evaluate(result).tolist(), expected)
+ result = ragged_array_ops.gather_nd(params, indices)
+ self.assertRaggedEqual(result, expected)
- @test_util.run_deprecated_v1
def testRaggedGatherNdUnknownRankError(self):
- params = ragged.constant([['a', 'b'], ['c', 'd']])
+ if context.executing_eagerly():
+ return
+ params = ragged_factory_ops.constant([['a', 'b'], ['c', 'd']])
indices1 = array_ops.placeholder(dtypes.int32, shape=None)
indices2 = array_ops.placeholder(dtypes.int32, shape=[None])
with self.assertRaisesRegexp(ValueError,
'indices.rank be statically known.'):
- ragged.gather_nd(params, indices1)
+ ragged_array_ops.gather_nd(params, indices1)
with self.assertRaisesRegexp(
ValueError, r'indices.shape\[-1\] must be statically known.'):
- ragged.gather_nd(params, indices2)
+ ragged_array_ops.gather_nd(params, indices2)
@parameterized.parameters([
dict(
params=['a'],
indices=0,
- message='Shape must be at least rank 1 but is rank 0'
- " for 'GatherNd'"),
+ error=(ValueError, errors.InvalidArgumentError)),
dict(
- params=ragged.constant_value([['a']]),
+ params=ragged_factory_ops.constant_value([['a']]),
indices=0,
message='indices.rank must be at least 1.'),
dict(
params=['a', 'b', 'c'],
- indices=ragged.constant([[0]]),
+ indices=ragged_factory_ops.constant_value([[0]]),
message='The innermost dimension of indices may not be ragged'),
])
- @test_util.run_deprecated_v1
def testRaggedGatherNdStaticError(self,
params,
indices,
- message,
+ message=None,
error=ValueError):
with self.assertRaisesRegexp(error, message):
- ragged.gather_nd(params, indices)
+ ragged_array_ops.gather_nd(params, indices)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_gather_op_test.py b/tensorflow/python/ops/ragged/ragged_gather_op_test.py
index 76c90cd..9914b56 100644
--- a/tensorflow/python/ops/ragged/ragged_gather_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_gather_op_test.py
@@ -12,142 +12,132 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.gather."""
+"""Tests for ragged_array_ops.gather."""
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 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 test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedTensorOpsTest(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedGatherOpTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testDocStringExamples(self):
params = constant_op.constant(['a', 'b', 'c', 'd', 'e'])
indices = constant_op.constant([3, 1, 2, 1, 0])
- ragged_params = ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
- ragged_indices = ragged.constant([[3, 1, 2], [1], [], [0]])
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, ragged_indices).eval().tolist(),
- [[b'd', b'b', b'c'], [b'b'], [], [b'a']])
- self.assertEqual(
- ragged.gather(ragged_params, indices).eval().tolist(),
- [[b'e'], [b'd'], [], [b'd'], [b'a', b'b', b'c']])
- self.assertEqual(
- ragged.gather(ragged_params, ragged_indices).eval().tolist(),
- [[[b'e'], [b'd'], []], [[b'd']], [], [[b'a', b'b', b'c']]])
+ ragged_params = ragged_factory_ops.constant([['a', 'b', 'c'], ['d'], [],
+ ['e']])
+ ragged_indices = ragged_factory_ops.constant([[3, 1, 2], [1], [], [0]])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, ragged_indices),
+ [[b'd', b'b', b'c'], [b'b'], [], [b'a']])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(ragged_params, indices),
+ [[b'e'], [b'd'], [], [b'd'], [b'a', b'b', b'c']])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(ragged_params, ragged_indices),
+ [[[b'e'], [b'd'], []], [[b'd']], [], [[b'a', b'b', b'c']]])
- @test_util.run_deprecated_v1
def testTensorParamsAndTensorIndices(self):
params = ['a', 'b', 'c', 'd', 'e']
indices = [2, 0, 2, 1]
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(),
- [b'c', b'a', b'c', b'b'])
- self.assertEqual(type(ragged.gather(params, indices)), ops.Tensor)
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices), [b'c', b'a', b'c', b'b'])
+ self.assertIsInstance(ragged_array_ops.gather(params, indices), ops.Tensor)
- @test_util.run_deprecated_v1
def testRaggedParamsAndTensorIndices(self):
- params = ragged.constant([['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']])
+ params = ragged_factory_ops.constant([['a', 'b'], ['c', 'd', 'e'], ['f'],
+ [], ['g']])
indices = [2, 0, 2, 1]
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(),
- [[b'f'], [b'a', b'b'], [b'f'], [b'c', b'd', b'e']])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices),
+ [[b'f'], [b'a', b'b'], [b'f'], [b'c', b'd', b'e']])
- @test_util.run_deprecated_v1
def testTensorParamsAndRaggedIndices(self):
params = ['a', 'b', 'c', 'd', 'e']
- indices = ragged.constant([[2, 1], [1, 2, 0], [3]])
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(),
- [[b'c', b'b'], [b'b', b'c', b'a'], [b'd']])
+ indices = ragged_factory_ops.constant([[2, 1], [1, 2, 0], [3]])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices),
+ [[b'c', b'b'], [b'b', b'c', b'a'], [b'd']])
- @test_util.run_deprecated_v1
def testRaggedParamsAndRaggedIndices(self):
- params = ragged.constant([['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']])
- indices = ragged.constant([[2, 1], [1, 2, 0], [3]])
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(),
- [[[b'f'], [b'c', b'd', b'e']], # [[p[2], p[1] ],
- [[b'c', b'd', b'e'], [b'f'], [b'a', b'b']], # [p[1], p[2], p[0]],
- [[]]] # [p[3] ]]
- ) # pyformat: disable
+ params = ragged_factory_ops.constant([['a', 'b'], ['c', 'd', 'e'], ['f'],
+ [], ['g']])
+ indices = ragged_factory_ops.constant([[2, 1], [1, 2, 0], [3]])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices),
+ [[[b'f'], [b'c', b'd', b'e']], # [[p[2], p[1] ],
+ [[b'c', b'd', b'e'], [b'f'], [b'a', b'b']], # [p[1], p[2], p[0]],
+ [[]]] # [p[3] ]]
+ ) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedParamsAndScalarIndices(self):
- params = ragged.constant([['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']])
+ params = ragged_factory_ops.constant([['a', 'b'], ['c', 'd', 'e'], ['f'],
+ [], ['g']])
indices = 1
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(), [b'c', b'd', b'e'])
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices), [b'c', b'd', b'e'])
- @test_util.run_deprecated_v1
def test3DRaggedParamsAnd2DTensorIndices(self):
- params = ragged.constant([[['a', 'b'], []], [['c', 'd'], ['e'], ['f']],
- [['g']]])
+ params = ragged_factory_ops.constant([[['a', 'b'], []],
+ [['c', 'd'], ['e'], ['f']], [['g']]])
indices = [[1, 2], [0, 1], [2, 2]]
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(),
- [[[[b'c', b'd'], [b'e'], [b'f']], [[b'g']]], # [[p1, p2],
- [[[b'a', b'b'], []], [[b'c', b'd'], [b'e'], [b'f']]], # [p0, p1],
- [[[b'g']], [[b'g']]]] # [p2, p2]]
- ) # pyformat: disable
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices),
+ [[[[b'c', b'd'], [b'e'], [b'f']], [[b'g']]], # [[p1, p2],
+ [[[b'a', b'b'], []], [[b'c', b'd'], [b'e'], [b'f']]], # [p0, p1],
+ [[[b'g']], [[b'g']]]] # [p2, p2]]
+ ) # pyformat: disable
- @test_util.run_deprecated_v1
def testTensorParamsAnd4DRaggedIndices(self):
- indices = ragged.constant(
+ indices = ragged_factory_ops.constant(
[[[[3, 4], [0, 6]], []], [[[2, 1], [1, 0]], [[2, 5]], [[2, 3]]],
[[[1, 0]]]], # pyformat: disable
ragged_rank=2,
inner_shape=(2,))
params = ['a', 'b', 'c', 'd', 'e', 'f', 'g']
- with self.test_session():
- self.assertEqual(
- ragged.gather(params, indices).eval().tolist(),
- [[[[b'd', b'e'], [b'a', b'g']], []],
- [[[b'c', b'b'], [b'b', b'a']], [[b'c', b'f']], [[b'c', b'd']]],
- [[[b'b', b'a']]]]) # pyformat: disable
+ self.assertRaggedEqual(
+ ragged_array_ops.gather(params, indices),
+ [[[[b'd', b'e'], [b'a', b'g']], []],
+ [[[b'c', b'b'], [b'b', b'a']], [[b'c', b'f']], [[b'c', b'd']]],
+ [[[b'b', b'a']]]]) # pyformat: disable
- @test_util.run_deprecated_v1
def testOutOfBoundsError(self):
tensor_params = ['a', 'b', 'c']
tensor_indices = [0, 1, 2]
- ragged_params = ragged.constant([['a', 'b'], ['c']])
- ragged_indices = ragged.constant([[0, 3]])
- with self.test_session():
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'indices\[1\] = 3 is not in \[0, 3\)',
- ragged.gather(tensor_params, ragged_indices).eval)
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'indices\[2\] = 2 is not in \[0, 2\)',
- ragged.gather(ragged_params, tensor_indices).eval)
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'indices\[1\] = 3 is not in \[0, 2\)',
- ragged.gather(ragged_params, ragged_indices).eval)
+ ragged_params = ragged_factory_ops.constant([['a', 'b'], ['c']])
+ ragged_indices = ragged_factory_ops.constant([[0, 3]])
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'indices\[1\] = 3 is not in \[0, 3\)'):
+ self.evaluate(ragged_array_ops.gather(tensor_params, ragged_indices))
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'indices\[2\] = 2 is not in \[0, 2\)'):
+ self.evaluate(ragged_array_ops.gather(ragged_params, tensor_indices))
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'indices\[1\] = 3 is not in \[0, 2\)'):
+ self.evaluate(ragged_array_ops.gather(ragged_params, ragged_indices))
- @test_util.run_deprecated_v1
def testUnknownIndicesRankError(self):
- params = ragged.constant([], ragged_rank=1)
+ if context.executing_eagerly():
+ return
+ params = ragged_factory_ops.constant([], ragged_rank=1)
indices = constant_op.constant([0], dtype=dtypes.int64)
indices = array_ops.placeholder_with_default(indices, None)
self.assertRaisesRegexp(ValueError,
r'indices\.shape\.ndims must be known statically',
- ragged.gather, params, indices)
+ ragged_array_ops.gather, params, indices)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_getitem.py b/tensorflow/python/ops/ragged/ragged_getitem.py
index 9821695..001a400 100644
--- a/tensorflow/python/ops/ragged/ragged_getitem.py
+++ b/tensorflow/python/ops/ragged/ragged_getitem.py
@@ -24,7 +24,6 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.ragged import ragged_array_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_math_ops
from tensorflow.python.ops.ragged import ragged_tensor
@@ -39,7 +38,7 @@
IndexError; (2) use a default value; or (3) skip that value and return a
tensor with fewer rows than we started with. Following the guiding
principles of Python ("In the face of ambiguity, refuse the temptation to
- guess" <go/pep20>), we simply disallow this operation.
+ guess"), we simply disallow this operation.
Any dimensions added by `array_ops.newaxis` will be ragged if the following
dimension is ragged.
@@ -137,8 +136,8 @@
if row_key is array_ops.newaxis:
inner_rt = _ragged_getitem(rt_input, inner_keys)
nsplits = array_ops.shape(inner_rt.row_splits, out_type=dtypes.int64)[0]
- return ragged_factory_ops.from_row_splits(inner_rt,
- array_ops.stack([0, nsplits - 1]))
+ return ragged_tensor.RaggedTensor.from_row_splits(
+ inner_rt, array_ops.stack([0, nsplits - 1]))
# Slicing a range of rows: first slice the outer dimension, and then
# call `_ragged_getitem_inner_dimensions` to handle the inner keys.
@@ -184,7 +183,7 @@
axis=0)
values_start = new_splits[0]
values_limit = new_splits[-1]
- return ragged_factory_ops.from_row_splits(
+ return ragged_tensor.RaggedTensor.from_row_splits(
rt_input.values[values_start:values_limit], new_splits - values_start)
# If there is a slice step (aka a strided slice), then use ragged_gather to
@@ -225,7 +224,8 @@
if column_key is array_ops.newaxis:
inner_rt = _ragged_getitem_inner_dimensions(rt_input, key_list[1:])
nsplits = array_ops.shape(inner_rt.row_splits, out_type=dtypes.int64)[0]
- return ragged_factory_ops.from_row_splits(inner_rt, math_ops.range(nsplits))
+ return ragged_tensor.RaggedTensor.from_row_splits(inner_rt,
+ math_ops.range(nsplits))
# Slicing a range of columns in a ragged inner dimension. We use a
# recursive call to process the values, and then assemble a RaggedTensor
@@ -239,7 +239,7 @@
else:
# Nontrivial slice: use ragged_gather to extract the indicated slice as
# a new RaggedTensor (inner_rt), and then recursively process its values.
- # The splits can be taken from ragged.row_splits(inner_rt).
+ # The splits can be taken from inner_rt.row_splits().
inner_rt_starts = rt_input.row_splits[:-1]
inner_rt_limits = rt_input.row_splits[1:]
if column_key.start is not None and column_key.start != 0:
diff --git a/tensorflow/python/ops/ragged/ragged_map_flat_values_op_test.py b/tensorflow/python/ops/ragged/ragged_map_flat_values_op_test.py
new file mode 100644
index 0000000..e9a7cdf
--- /dev/null
+++ b/tensorflow/python/ops/ragged/ragged_map_flat_values_op_test.py
@@ -0,0 +1,207 @@
+# 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 ragged_functional_ops.map_flat_values."""
+
+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 errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedMapInnerValuesOpTest(ragged_test_util.RaggedTensorTestCase):
+
+ def assertRaggedMapInnerValuesReturns(self,
+ op,
+ expected,
+ args=(),
+ kwargs=None):
+ kwargs = kwargs or {}
+ result = ragged_functional_ops.map_flat_values(op, *args, **kwargs)
+ self.assertRaggedEqual(result, expected)
+
+ def testDocStringExamples(self):
+ """Test the examples in apply_op_to_ragged_values.__doc__."""
+ rt = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5], [6]])
+ v1 = ragged_functional_ops.map_flat_values(array_ops.ones_like, rt)
+ v2 = ragged_functional_ops.map_flat_values(math_ops.multiply, rt, rt)
+ v3 = ragged_functional_ops.map_flat_values(math_ops.add, rt, 5)
+ self.assertRaggedEqual(v1, [[1, 1, 1], [], [1, 1], [1]])
+ self.assertRaggedEqual(v2, [[1, 4, 9], [], [16, 25], [36]])
+ self.assertRaggedEqual(v3, [[6, 7, 8], [], [9, 10], [11]])
+
+ def testOpWithSingleRaggedTensorArg(self):
+ tensor = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=array_ops.zeros_like,
+ args=(tensor,),
+ expected=[[0, 0, 0], [], [0, 0]])
+
+ def testOpWithTwoRaggedTensorArgs(self):
+ x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+ y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply, args=(x, y), expected=[[3, 2, 12], [], [4, 25]])
+
+ def testOpWithRaggedTensorAndScalarArgs(self):
+ y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply, args=(5, y), expected=[[5, 10, 15], [], [20, 25]])
+
+ def testOpWithThreeRaggedTensorArgs(self):
+ condition = ragged_factory_ops.constant(
+ [[True, True, False], [], [True, False]]) # pyformat: disable
+ x = ragged_factory_ops.constant([['a', 'b', 'c'], [], ['d', 'e']])
+ y = ragged_factory_ops.constant([['A', 'B', 'C'], [], ['D', 'E']])
+ self.assertRaggedMapInnerValuesReturns(
+ op=array_ops.where,
+ args=(condition, x, y),
+ expected=[[b'a', b'b', b'C'], [], [b'd', b'E']])
+
+ def testOpWithRaggedTensorListArg(self):
+ x = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+ y = ragged_factory_ops.constant([[10, 20, 30], [], [40, 50]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.add_n,
+ args=([x, y, x],),
+ expected=[[12, 24, 36], [], [48, 60]])
+
+ def testOpWithKeywordArgs(self):
+ x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+ y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply,
+ kwargs=dict(x=x, y=y),
+ expected=[[3, 2, 12], [], [4, 25]])
+
+ def testOpWithMixedPositionalAndKeywordArgs(self):
+ x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+ y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply,
+ args=(x,),
+ kwargs=dict(y=y),
+ expected=[[3, 2, 12], [], [4, 25]])
+
+ def testNonElementWiseOp(self):
+ x = ragged_factory_ops.constant(
+ [[[3, 1, 4], [1, 5, 9], [2, 6, 5]], [], [[3, 5, 8], [9, 7, 9]]],
+ ragged_rank=1)
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.reduce_sum,
+ kwargs={
+ 'input_tensor': x,
+ 'axis': 1,
+ },
+ expected=[[8, 15, 13], [], [16, 25]])
+
+ def testOpWithRaggedRankGreaterThanOne(self):
+ # ragged_rank=0
+ x0 = [3, 1, 4, 1, 5, 9, 2, 6, 5]
+ y0 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+ self.assertRaggedEqual(
+ math_ops.multiply(x0, y0), [3, 2, 12, 4, 25, 54, 14, 48, 45])
+
+ # ragged_rank=1
+ x1 = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5], [9, 2], [6, 5]])
+ y1 = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5], [6, 7], [8, 9]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply,
+ args=(x1, y1),
+ expected=[[3, 2, 12], [], [4, 25], [54, 14], [48, 45]])
+
+ # ragged_rank=2
+ x2 = ragged_factory_ops.constant([[[3, 1, 4]], [], [[], [1, 5]],
+ [[9, 2], [6, 5]]])
+ y2 = ragged_factory_ops.constant([[[1, 2, 3]], [], [[], [4, 5]],
+ [[6, 7], [8, 9]]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply,
+ args=(x2, y2),
+ expected=[[[3, 2, 12]], # row 0
+ [], # row 1
+ [[], [4, 25]], # row 2
+ [[54, 14], [48, 45]] # row 3
+ ]) # pyformat: disable
+
+ # ragged_rank=3
+ x3 = ragged_factory_ops.constant([[[[3, 1, 4]], []], [], [[[], [1, 5]]],
+ [[[9, 2], [6, 5]]]])
+ y3 = ragged_factory_ops.constant([[[[1, 2, 3]], []], [], [[[], [4, 5]]],
+ [[[6, 7], [8, 9]]]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply,
+ args=(x3, y3),
+ expected=[
+ [[[3, 2, 12]], []], # row 0
+ [], # row 1
+ [[[], [4, 25]]], # row 2
+ [[[54, 14], [48, 45]]] # row 3
+ ]) # pyformat: disable
+
+ def testOpWithRaggedRankThree(self):
+ x = ragged_factory_ops.constant([[[3, 1, 4]], [], [[], [1, 5]]])
+ y = ragged_factory_ops.constant([[[1, 2, 3]], [], [[], [4, 5]]])
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply,
+ args=(x, y),
+ expected=[[[3, 2, 12]], [], [[], [4, 25]]])
+
+ def testOpWithInnerValuesOnly(self):
+ x = constant_op.constant([[1, 2], [3, 4], [5, 6]])
+ y = constant_op.constant(2)
+ self.assertRaggedMapInnerValuesReturns(
+ op=math_ops.multiply, args=(x, y), expected=[[2, 4], [6, 8], [10, 12]])
+
+ def testRaggedTensorSplitsRaggedRankMismatchError(self):
+ x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+ y = ragged_factory_ops.constant([[[3, 1, 4], []], [], [[1, 5]]])
+ self.assertRaisesRegexp(
+ ValueError, r'Inputs must have identical ragged splits.*',
+ ragged_functional_ops.map_flat_values, math_ops.add, x, y)
+
+ def testRaggedTensorSplitsValueMismatchError(self):
+ x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+ y = ragged_factory_ops.constant([[1], [2, 3], [4, 5]])
+ self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'Inputs must have identical ragged splits.*',
+ ragged_functional_ops.map_flat_values, math_ops.add,
+ x, y)
+
+ def testRaggedTensorSplitsMismatchErrorAtRuntime(self):
+ splits1 = array_ops.placeholder_with_default(
+ constant_op.constant([0, 3, 3, 5], dtypes.int64), None)
+ splits2 = array_ops.placeholder_with_default(
+ constant_op.constant([0, 1, 3, 5], dtypes.int64), None)
+ x = ragged_tensor.RaggedTensor.from_row_splits([3, 1, 4, 1, 5], splits1)
+ y = ragged_tensor.RaggedTensor.from_row_splits([1, 2, 3, 4, 5], splits2)
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'.*Inputs must have identical ragged splits'):
+ self.evaluate(ragged_functional_ops.map_flat_values(math_ops.add, x, y))
+
+
+if __name__ == '__main__':
+ googletest.main()
diff --git a/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py b/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py
index 7a8603c..1520640 100644
--- a/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_map_fn_op_test.py
@@ -12,24 +12,34 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.map_fn."""
+"""Tests for ragged_map_ops.map_fn."""
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.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.keras import backend
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops as mo
-from tensorflow.python.ops import ragged
from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_map_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedMapOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedMapOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
+
@parameterized.parameters([
# The following test sets map over a RaggedTensor and apply a
# transformation that returns with shape:
@@ -52,61 +62,62 @@
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[2, 6], [4.5, 9], [6.5, 13]],
dtype=dtypes.float32,
+ expected_ragged_rank=0,
),
# [d1, (d2)] -> [d1, (d2)]
dict(
- fn=lambda x: x+1,
+ fn=lambda x: x + np.int64(1),
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[2, 3, 4], [5, 6], [7, 8]],
dtype=dtypes.int64,
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=1),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=1),
),
# [d1, (d2), d3] -> [d1, (d2), d3]
dict(
- fn=lambda x: x+1,
+ fn=lambda x: x + np.int64(1),
elems=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
elems_ragged_rank=1,
expected_ragged_rank=1,
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=1),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=1),
expected_output=[[[2, 3], [4, 5]], [], [[6, 7], [8, 9], [10, 1]]],
),
# [d1, (d2)] -> [d1, (d2), (d3)]
dict(
- fn=lambda x: ragged.from_row_starts(x, [0]),
+ fn=lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0]),
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[[1, 2, 3]], [[4, 5]], [[6, 7]]],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=2),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=2),
),
# [d1, (d2), (d3)] -> [d1, (d2), (d3)]
dict(
- fn=lambda x: ragged.map_inner_values(mo.add, x, 1),
+ fn=lambda x: ragged_functional_ops.map_flat_values(mo.add, x, 1),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[[2, 3, 4]], [[5, 6], [7, 8]]],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=2),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=2),
),
# [d1, (d2), (d3)] -> [d1, (d2)]
dict(
- fn=lambda x: ragged.reduce_sum(x, axis=1),
+ fn=lambda x: ragged_math_ops.reduce_sum(x, axis=1),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[6], [9, 13]],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=1),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=1),
),
# [d1, (d2), (d3)] -> [d1, (d3)]
dict(
- fn=lambda x: ragged.reduce_sum(x, axis=0),
+ fn=lambda x: ragged_math_ops.reduce_sum(x, axis=0),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[1, 2, 3], [10, 12]],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=1),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=1),
),
# [d1, (d2), (d3)] -> [d1]
dict(
- fn=ragged.reduce_sum,
+ fn=ragged_math_ops.reduce_sum,
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[6, 22],
result_dtype=dtypes.int64,
@@ -116,31 +127,29 @@
fn=mo.range,
elems=[4, 0, 2],
expected_output=[[0, 1, 2, 3], [], [0, 1]],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=1),
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=1),
),
# [d1] -> [d1, (d2), (d3)]
dict(
- fn=lambda x: ragged.range(mo.range(x)),
+ fn=lambda x: ragged_math_ops.range(mo.range(x)),
elems=[5, 0, 3],
- expected_output=[
- [[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]], [], [[], [0], [0, 1]]
- ],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=2),
+ expected_output=[[[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]], [],
+ [[], [0], [0, 1]]],
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=2),
),
# [d1, (d2), (d3), (d4a), (d5)] -> [d1, (d2), (d3), (d4b), (d5)]
dict(
- fn=lambda x: ragged.add(x, 1),
+ fn=lambda x: x + np.int64(1),
elems=[[[[[1, 2, 3]], [[4], [5]]]], [[[[6, 7]]], [[[8], []]]]],
- expected_output=[[[[[2, 3, 4]], [[5], [6]]]],
- [[[[7, 8]]], [[[9], []]]]],
- result_dtype=ragged.RaggedTensorType(dtype=dtypes.int64,
- ragged_rank=4),
+ expected_output=[[[[[2, 3, 4]], [[5], [6]]]], [[[[7, 8]]], [[[9],
+ []]]]],
+ result_dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=4),
),
])
- @test_util.run_deprecated_v1
def testRaggedMap(
self,
fn,
@@ -153,135 +162,135 @@
result_dtype=None,
infer_shape=False,
):
- elems = ragged.constant(elems, dtype, elems_ragged_rank)
- output = ragged.map_fn(
+ elems = ragged_factory_ops.constant(elems, dtype, elems_ragged_rank)
+ output = ragged_map_ops.map_fn(
fn=fn, elems=elems, dtype=result_dtype, infer_shape=infer_shape)
- expected_rt = ragged.constant(
+ expected_rt = ragged_factory_ops.constant(
expected_output, ragged_rank=expected_ragged_rank)
- with self.test_session():
- if ragged.is_ragged(expected_output):
- self.assertEqual(output.ragged_rank, expected_rt.ragged_rank)
- output_values = self.evaluate(output)
- self.assertAllEqual(expected_output, output_values.tolist())
+ self.assertRaggedEqual(expected_rt, output)
- @test_util.run_deprecated_v1
def testRaggedMapOnStructure(self):
- batman = ragged.constant([[1, 2, 3], [4], [5, 6, 7]])
+ batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
# [[10, 20, 30], [40], [50, 60, 70]]
- robin = ragged.map_inner_values(mo.multiply, batman, 10)
+ robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
features = {'batman': batman, 'robin': robin}
def _reduce_sum_from_all(f):
return mo.reduce_sum(f['batman']) + mo.reduce_sum(f['robin'])
- output = ragged.map_fn(
+ output = ragged_map_ops.map_fn(
fn=_reduce_sum_from_all,
elems=features,
dtype=dtypes.int32,
)
- with self.test_session():
- self.assertAllEqual(output.eval().tolist(), [66, 44, 198])
+ self.assertRaggedEqual(output, [66, 44, 198])
# Test mapping over a dict of RTs can produce a dict of RTs.
- @test_util.run_deprecated_v1
def testRaggedMapOnStructure_RaggedOutputs(self):
- batman = ragged.constant([[1, 2, 3], [4], [5, 6, 7]])
+ batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
# [[10, 20, 30], [40], [50, 60, 70]]
- robin = ragged.map_inner_values(mo.multiply, batman, 10)
+ robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
features = {'batman': batman, 'robin': robin}
def _increment(f):
return {
- 'batman': ragged.add(f['batman'], 1),
- 'robin': ragged.add(f['robin'], 1),
+ 'batman': f['batman'] + 1,
+ 'robin': f['robin'] + 1,
}
- output = ragged.map_fn(
+ output = ragged_map_ops.map_fn(
fn=_increment,
elems=features,
infer_shape=False,
dtype={
'batman':
- ragged.RaggedTensorType(dtype=dtypes.int32, ragged_rank=1),
+ ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int32, ragged_rank=1),
'robin':
- ragged.RaggedTensorType(dtype=dtypes.int32, ragged_rank=1)
+ ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int32, ragged_rank=1)
},
)
- with self.test_session():
- self.assertAllEqual(output['batman'].eval().tolist(),
- [[2, 3, 4], [5], [6, 7, 8]])
- self.assertAllEqual(output['robin'].eval().tolist(),
- [[11, 21, 31], [41], [51, 61, 71]])
+ self.assertRaggedEqual(output['batman'], [[2, 3, 4], [5], [6, 7, 8]])
+ self.assertRaggedEqual(output['robin'], [[11, 21, 31], [41], [51, 61, 71]])
- @test_util.run_deprecated_v1
def testZip(self):
- x = ragged.constant([[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]],
- dtypes.int64)
- y = array_ops.expand_dims(
- mo.range(ragged.nrows(x), dtype=dtypes.int64), axis=1)
+ x = ragged_factory_ops.constant(
+ [[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]], dtypes.int64)
+ y = array_ops.expand_dims(mo.range(x.nrows(), dtype=dtypes.int64), axis=1)
def _zip(foo):
y_val, x_val = foo
bar = backend.tile(y_val, array_ops.shape(x_val))
return array_ops.stack([bar, x_val], axis=1)
- output = ragged.map_fn(
+ output = ragged_map_ops.map_fn(
_zip, (y, x),
- dtype=ragged.RaggedTensorType(dtype=dtypes.int64, ragged_rank=1),
+ dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64, ragged_rank=1),
infer_shape=False)
- with self.test_session():
- result = self.evaluate(output).tolist()
- self.assertAllEqual(
- result, [[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]],
- [[3, 70]], [[4, 80], [4, 90], [4, 100]]])
+ self.assertRaggedEqual(
+ output, [[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]],
+ [[3, 70]], [[4, 80], [4, 90], [4, 100]]])
- @test_util.run_deprecated_v1
def testBatchGather(self):
- tokens = ragged.constant([['hello', '.', 'there'], ['merhaba'],
- ['bonjour', '.', 'ca va', '?']])
- indices = ragged.constant([[0, 2], [0], [0, 2]])
+ tokens = ragged_factory_ops.constant([['hello', '.', 'there'], ['merhaba'],
+ ['bonjour', '.', 'ca va', '?']])
+ indices = ragged_factory_ops.constant([[0, 2], [0], [0, 2]])
def gather(x):
tokens_val, indices_val = x
return array_ops.gather(tokens_val, indices_val)
data = tokens, indices
- out = ragged.map_fn(
+ out = ragged_map_ops.map_fn(
gather,
data,
- dtype=ragged.RaggedTensorType(dtype=dtypes.string, ragged_rank=1),
+ dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.string, ragged_rank=1),
infer_shape=False)
- with self.test_session():
- self.assertAllEqual(
- self.evaluate(out).tolist(),
- [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
+ self.assertRaggedEqual(
+ out, [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
def testMismatchRaggedRank(self):
- elems = ragged.constant([[[1, 2, 3]], [[4, 5], [6, 7]]])
- fn = lambda x: ragged.reduce_sum(x, axis=0)
+ elems = ragged_factory_ops.constant([[[1, 2, 3]], [[4, 5], [6, 7]]])
+ fn = lambda x: ragged_math_ops.reduce_sum(x, axis=0)
with self.assertRaisesWithLiteralMatch(
ValueError, r'The declared ragged rank (23) mismatches the result (1)'):
- _ = ragged.map_fn(
+ _ = ragged_map_ops.map_fn(
fn,
elems,
- dtype=ragged.RaggedTensorType(dtype=dtypes.int64, ragged_rank=23))
+ dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=23))
def testMismatchRaggedRank2(self):
- elems = ragged.constant([[1, 2, 3], [4, 5], [6, 7]])
- fn = lambda x: ragged.from_row_starts(x, [0])
+ elems = ragged_factory_ops.constant([[1, 2, 3], [4, 5], [6, 7]])
+ fn = lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0])
with self.assertRaisesWithLiteralMatch(
ValueError, r'The declared ragged rank (10) mismatches the result (1)'):
- _ = ragged.map_fn(
+ _ = ragged_map_ops.map_fn(
fn,
elems,
- dtype=ragged.RaggedTensorType(dtype=dtypes.int64, ragged_rank=10))
+ dtype=ragged_tensor.RaggedTensorType(
+ dtype=dtypes.int64, ragged_rank=10))
+
+ def testMapOnSparseTensor(self):
+ s = sparse_tensor.SparseTensor(
+ indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
+ values=[0, 5, 0, 4],
+ dense_shape=[2, 2],
+ )
+ t2 = ragged_tensor.RaggedTensor.from_sparse(s)
+ id_t2 = ragged_map_ops.map_fn(
+ lambda x: x, t2,
+ )
+ self.assertRaggedEqual(id_t2, [[0, 5], [0, 4]])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_map_inner_values_op_test.py b/tensorflow/python/ops/ragged/ragged_map_inner_values_op_test.py
deleted file mode 100644
index b5802cb..0000000
--- a/tensorflow/python/ops/ragged/ragged_map_inner_values_op_test.py
+++ /dev/null
@@ -1,223 +0,0 @@
-# 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 ragged.map_inner_values."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
-from tensorflow.python.framework import test_util
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import ragged
-from tensorflow.python.platform import googletest
-
-
-class RaggedMapInnerValuesOpTest(test_util.TensorFlowTestCase,
- parameterized.TestCase):
-
- def assertRaggedMapInnerValuesReturns(self,
- op,
- expected,
- args=(),
- kwargs=None):
- kwargs = kwargs or {}
- result = ragged.map_inner_values(op, *args, **kwargs)
- with self.test_session():
- self.assertEqual(result.eval().tolist(), expected)
-
- @test_util.run_deprecated_v1
- def testDocStringExamples(self):
- """Test the examples in apply_op_to_ragged_values.__doc__."""
- rt = ragged.constant([[1, 2, 3], [], [4, 5], [6]])
- v1 = ragged.map_inner_values(array_ops.ones_like, rt)
- v2 = ragged.map_inner_values(math_ops.multiply, rt, rt)
- v3 = ragged.map_inner_values(math_ops.add, rt, 5)
- with self.test_session():
- self.assertEqual(v1.eval().tolist(), [[1, 1, 1], [], [1, 1], [1]])
- self.assertEqual(v2.eval().tolist(), [[1, 4, 9], [], [16, 25], [36]])
- self.assertEqual(v3.eval().tolist(), [[6, 7, 8], [], [9, 10], [11]])
-
- @test_util.run_deprecated_v1
- def testOpWithSingleRaggedTensorArg(self):
- tensor = ragged.constant([[1, 2, 3], [], [4, 5]])
- self.assertRaggedMapInnerValuesReturns(
- op=array_ops.zeros_like,
- args=(tensor,),
- expected=[[0, 0, 0], [], [0, 0]])
-
- @test_util.run_deprecated_v1
- def testOpWithTwoRaggedTensorArgs(self):
- x = ragged.constant([[3, 1, 4], [], [1, 5]])
- y = ragged.constant([[1, 2, 3], [], [4, 5]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply, args=(x, y), expected=[[3, 2, 12], [], [4, 25]])
-
- @test_util.run_deprecated_v1
- def testOpWithRaggedTensorAndScalarArgs(self):
- y = ragged.constant([[1, 2, 3], [], [4, 5]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply, args=(5, y), expected=[[5, 10, 15], [], [20, 25]])
-
- @test_util.run_deprecated_v1
- def testOpWithThreeRaggedTensorArgs(self):
- condition = ragged.constant(
- [[True, True, False], [], [True, False]]) # pyformat: disable
- x = ragged.constant([['a', 'b', 'c'], [], ['d', 'e']])
- y = ragged.constant([['A', 'B', 'C'], [], ['D', 'E']])
- self.assertRaggedMapInnerValuesReturns(
- op=array_ops.where,
- args=(condition, x, y),
- expected=[[b'a', b'b', b'C'], [], [b'd', b'E']])
-
- @test_util.run_deprecated_v1
- def testOpWithRaggedTensorListArg(self):
- x = ragged.constant([[1, 2, 3], [], [4, 5]])
- y = ragged.constant([[10, 20, 30], [], [40, 50]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.add_n,
- args=([x, y, x],),
- expected=[[12, 24, 36], [], [48, 60]])
-
- @test_util.run_deprecated_v1
- def testOpWithKeywordArgs(self):
- x = ragged.constant([[3, 1, 4], [], [1, 5]])
- y = ragged.constant([[1, 2, 3], [], [4, 5]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply,
- kwargs=dict(x=x, y=y),
- expected=[[3, 2, 12], [], [4, 25]])
-
- @test_util.run_deprecated_v1
- def testOpWithMixedPositionalAndKeywordArgs(self):
- x = ragged.constant([[3, 1, 4], [], [1, 5]])
- y = ragged.constant([[1, 2, 3], [], [4, 5]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply,
- args=(x,),
- kwargs=dict(y=y),
- expected=[[3, 2, 12], [], [4, 25]])
-
- @test_util.run_deprecated_v1
- def testNonElementWiseOp(self):
- x = ragged.constant(
- [[[3, 1, 4], [1, 5, 9], [2, 6, 5]], [], [[3, 5, 8], [9, 7, 9]]],
- ragged_rank=1)
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.reduce_sum,
- kwargs={
- 'input_tensor': x,
- 'axis': 1,
- },
- expected=[[8, 15, 13], [], [16, 25]])
-
- @test_util.run_deprecated_v1
- def testOpWithRaggedRankGreaterThanOne(self):
- # ragged_rank=0
- x0 = [3, 1, 4, 1, 5, 9, 2, 6, 5]
- y0 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
- with self.test_session():
- self.assertEqual(
- math_ops.multiply(x0, y0).eval().tolist(),
- [3, 2, 12, 4, 25, 54, 14, 48, 45])
-
- # ragged_rank=1
- x1 = ragged.constant([[3, 1, 4], [], [1, 5], [9, 2], [6, 5]])
- y1 = ragged.constant([[1, 2, 3], [], [4, 5], [6, 7], [8, 9]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply,
- args=(x1, y1),
- expected=[[3, 2, 12], [], [4, 25], [54, 14], [48, 45]])
-
- # ragged_rank=2
- x2 = ragged.constant([[[3, 1, 4]], [], [[], [1, 5]], [[9, 2], [6, 5]]])
- y2 = ragged.constant([[[1, 2, 3]], [], [[], [4, 5]], [[6, 7], [8, 9]]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply,
- args=(x2, y2),
- expected=[[[3, 2, 12]], # row 0
- [], # row 1
- [[], [4, 25]], # row 2
- [[54, 14], [48, 45]] # row 3
- ]) # pyformat: disable
-
- # ragged_rank=3
- x3 = ragged.constant([[[[3, 1, 4]], []], [], [[[], [1, 5]]],
- [[[9, 2], [6, 5]]]])
- y3 = ragged.constant([[[[1, 2, 3]], []], [], [[[], [4, 5]]],
- [[[6, 7], [8, 9]]]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply,
- args=(x3, y3),
- expected=[
- [[[3, 2, 12]], []], # row 0
- [], # row 1
- [[[], [4, 25]]], # row 2
- [[[54, 14], [48, 45]]] # row 3
- ]) # pyformat: disable
-
- @test_util.run_deprecated_v1
- def testOpWithRaggedRankThree(self):
- x = ragged.constant([[[3, 1, 4]], [], [[], [1, 5]]])
- y = ragged.constant([[[1, 2, 3]], [], [[], [4, 5]]])
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply,
- args=(x, y),
- expected=[[[3, 2, 12]], [], [[], [4, 25]]])
-
- @test_util.run_deprecated_v1
- def testOpWithInnerValuesOnly(self):
- x = constant_op.constant([[1, 2], [3, 4], [5, 6]])
- y = constant_op.constant(2)
- self.assertRaggedMapInnerValuesReturns(
- op=math_ops.multiply, args=(x, y), expected=[[2, 4], [6, 8], [10, 12]])
-
- def testRaggedTensorSplitsRaggedRankMismatchError(self):
- x = ragged.constant([[3, 1, 4], [], [1, 5]])
- y = ragged.constant([[[3, 1, 4], []], [], [[1, 5]]])
- self.assertRaisesRegexp(ValueError,
- r'Inputs must have identical ragged splits.*',
- ragged.map_inner_values, math_ops.add, x, y)
-
- def testRaggedTensorSplitsValueMismatchError(self):
- x = ragged.constant([[3, 1, 4], [], [1, 5]])
- y = ragged.constant([[1], [2, 3], [4, 5]])
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'Inputs must have identical ragged splits.*',
- ragged.map_inner_values, math_ops.add, x, y)
-
- @test_util.run_deprecated_v1
- def testRaggedTensorSplitsMismatchErrorAtRuntime(self):
- splits1 = array_ops.placeholder_with_default(
- constant_op.constant([0, 3, 3, 5], dtypes.int64), None)
- splits2 = array_ops.placeholder_with_default(
- constant_op.constant([0, 1, 3, 5], dtypes.int64), None)
- x = ragged.from_row_splits([3, 1, 4, 1, 5], splits1)
- y = ragged.from_row_splits([1, 2, 3, 4, 5], splits2)
- result = ragged.map_inner_values(math_ops.add, x, y)
- with self.test_session():
- self.assertRaisesRegexp(
- errors.InvalidArgumentError,
- r'\[Inputs must have identical ragged splits\] '
- r'\[Condition x == y did not hold element-wise:\].*', result.eval)
-
-
-if __name__ == '__main__':
- googletest.main()
diff --git a/tensorflow/python/ops/ragged/ragged_map_ops.py b/tensorflow/python/ops/ragged/ragged_map_ops.py
index fafa23b..fbe188b 100644
--- a/tensorflow/python/ops/ragged/ragged_map_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_map_ops.py
@@ -27,12 +27,12 @@
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.ops import array_ops
from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops import variable_scope as vs
-from tensorflow.python.ops.ragged import ragged_array_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import nest
@@ -216,8 +216,8 @@
varscope_caching_device_was_none = True
elems_flat = [
- ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- elem, name="elem") for elem in elems_flat
+ ragged_tensor.convert_to_tensor_or_ragged_tensor(elem, name="elem")
+ for elem in elems_flat
]
# We can either infer the output, or we can assume that it will be the same
@@ -226,7 +226,7 @@
# Find the number of iterations, n may be known statically.
if isinstance(elems_flat[0], ragged_tensor.RaggedTensor):
- n = ragged_array_ops.nrows(elems_flat[0], out_type=dtypes.int32)
+ n = elems_flat[0].nrows(out_type=dtypes.int32)
else:
static_shape = elems_flat[0].shape
if static_shape.ndims is not None and static_shape.ndims < 1:
@@ -236,8 +236,10 @@
else:
raise ValueError(
"elements in elems must be 1+ dimensional Tensors, not scalars")
- n = static_shape[0].value or array_ops.shape(elems_flat[0])[0]
+ n = (tensor_shape.dimension_value(static_shape[0]) or
+ array_ops.shape(elems_flat[0])[0])
+ n = math_ops.cast(n, dtype=dtypes.int32)
# Create a flat list of TAs.
# Flatten the dtype structure to a list.
@@ -254,7 +256,7 @@
for t in dtype_components_flat
]
- i = constant_op.constant(0)
+ i = constant_op.constant(0, dtype=dtypes.int32)
def compute(i, tas):
"""The loop body of map_fn.
@@ -334,7 +336,7 @@
class _RaggedTensorComponents(
collections.namedtuple(
"_RaggedTensorComponents",
- ["inner_values", "nested_row_lengths", "outer_row_length"])):
+ ["flat_values", "nested_row_lengths", "outer_row_length"])):
"""A namedtuple of components which represent a `RaggedTensor`.
_RaggedTensorComponents is a list of components which can be used to create a
@@ -344,7 +346,7 @@
The following are a list of components for a `RaggedTensor`:
- inner_values: The flat and inner values of a RaggedTensor. This could be
+ flat_values: The flat and inner values of a RaggedTensor. This could be
a `Tensor`, a `TensorArray`, or a data type.
nested_row_lengths: a tuple containing the row lengths of each rank. The
elements of the tuple could be `Tensor`s or `TensorArray`s.
@@ -357,12 +359,12 @@
def _concat_ragged_tensor_components(rt_ta):
- inner_values = rt_ta.inner_values.concat()
+ flat_values = rt_ta.flat_values.concat()
nested_row_lengths = tuple(
row_lengths_ta.concat() for row_lengths_ta in rt_ta.nested_row_lengths)
outer_row_length = rt_ta.outer_row_length.concat()
return _RaggedTensorComponents(
- inner_values=inner_values,
+ flat_values=flat_values,
nested_row_lengths=nested_row_lengths,
outer_row_length=outer_row_length)
@@ -374,17 +376,17 @@
# The three component pieces we need:
# - inner values
- inner_values = rt.inner_values
+ flat_values = rt.flat_values
# - row_splits of the RT
splits = rt.nested_row_splits
nested_row_lengths = tuple(split[1:] - split[:-1] for split in splits)
# - outer row length
- outer_row_length = array_ops.expand_dims(ragged_array_ops.nrows(rt), axis=0)
+ outer_row_length = array_ops.expand_dims(rt.nrows(), axis=0)
return _RaggedTensorComponents(
- inner_values=inner_values,
+ flat_values=flat_values,
nested_row_lengths=nested_row_lengths,
outer_row_length=outer_row_length,
)
@@ -395,11 +397,12 @@
if not isinstance(t, _RaggedTensorComponents):
return t
- values = t.inner_values
+ values = t.flat_values
nested_row_lengths = tuple(t.nested_row_lengths)
for nested_row_length in reversed(nested_row_lengths):
- values = ragged_factory_ops.from_row_lengths(values, nested_row_length)
- return ragged_factory_ops.from_row_lengths(values, t.outer_row_length)
+ values = ragged_tensor.RaggedTensor.from_row_lengths(
+ values, nested_row_length)
+ return ragged_tensor.RaggedTensor.from_row_lengths(values, t.outer_row_length)
def _maybe_decompose_dtype(d):
@@ -408,7 +411,7 @@
return d
result = _RaggedTensorComponents(
- inner_values=d.dtype,
+ flat_values=d.dtype,
nested_row_lengths=tuple(dtypes.int64 for i in range(d.ragged_rank - 1)),
outer_row_length=dtypes.int64,
)
@@ -435,10 +438,13 @@
"The declared ragged rank (%d) mismatches the result (1)" %
declared.ragged_rank)
- row_length = array_ops.expand_dims(
- ragged_array_ops.nrows(current), axis=0)
+ if isinstance(current, ragged_tensor.RaggedTensor):
+ nrows = current.nrows()
+ else:
+ nrows = array_ops.shape(current, out_type=dtypes.int64)[0]
+ row_length = array_ops.expand_dims(nrows, axis=0)
rt = _RaggedTensorComponents(
- inner_values=current,
+ flat_values=current,
nested_row_lengths=(),
outer_row_length=row_length)
yield rt
diff --git a/tensorflow/python/ops/ragged/ragged_math_ops.py b/tensorflow/python/ops/ragged/ragged_math_ops.py
index 857b8db..02e927b 100644
--- a/tensorflow/python/ops/ragged/ragged_math_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_math_ops.py
@@ -18,6 +18,8 @@
from __future__ import division
from __future__ import print_function
+import numpy as np
+
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
@@ -25,17 +27,18 @@
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import gen_ragged_math_ops
from tensorflow.python.ops import math_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_functional_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged import ragged_util
from tensorflow.python.ops.ragged import segment_id_ops
+from tensorflow.python.util.tf_export import tf_export
#===============================================================================
# ragged.range
#===============================================================================
# pylint: disable=redefined-builtin
+@tf_export('ragged.range')
def range(starts, limits=None, deltas=1, dtype=None, name=None):
"""Returns a `RaggedTensor` containing the specified sequences of numbers.
@@ -97,8 +100,8 @@
[dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64])
result = gen_ragged_math_ops.ragged_range(starts, limits, deltas, name=name)
- return ragged_factory_ops.from_row_splits(result.rt_dense_values,
- result.rt_nested_splits)
+ return ragged_tensor.RaggedTensor.from_row_splits(result.rt_dense_values,
+ result.rt_nested_splits)
def _infer_matching_dtype(tensors, dtype_hierarchy):
@@ -143,8 +146,11 @@
"""
-def _ragged_segment_aggregate(unsorted_segment_op, data, segment_ids,
- num_segments, name=None):
+def _ragged_segment_aggregate(unsorted_segment_op,
+ data,
+ segment_ids,
+ num_segments,
+ name=None):
"""Aggregates along segments of a RaggedTensor using `unsorted_segment_op`.
Returns a RaggedTensor `output` with `num_segments` rows, where the row
@@ -181,9 +187,8 @@
with ops.name_scope(name, 'RaggedSegment',
[data, segment_ids, num_segments]) as name:
- data = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- data, name='data')
- segment_ids = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
+ data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name='data')
+ segment_ids = ragged_tensor.convert_to_tensor_or_ragged_tensor(
segment_ids, name='segment_ids')
if ragged_tensor.is_ragged(segment_ids):
@@ -212,12 +217,11 @@
assert output_row_lengths.dtype == dtypes.int64
# Build the splits tensor for the output RaggedTensor.
- output_splits = array_ops.concat(
- [
- array_ops.zeros([1], dtypes.int64),
- math_ops.cumsum(output_row_lengths)
- ],
- axis=0)
+ output_splits = array_ops.concat([
+ array_ops.zeros([1], dtypes.int64),
+ math_ops.cumsum(output_row_lengths)
+ ],
+ axis=0)
# For each row in `data`, find the start & limit position where that row's
# values will be aggregated in output.values.
@@ -234,7 +238,8 @@
output_values = _ragged_segment_aggregate(unsorted_segment_op, data.values,
data_val_to_out_val_index,
output_splits[-1])
- return ragged_factory_ops.from_row_splits(output_values, output_splits)
+ return ragged_tensor.RaggedTensor.from_row_splits(output_values,
+ output_splits)
def segment_sum(data, segment_ids, num_segments, name=None):
@@ -266,28 +271,32 @@
def segment_mean(data, segment_ids, num_segments, name=None):
- # For docs, see: _RAGGED_SEGMENT_DOCSTRING
+ """For docs, see: _RAGGED_SEGMENT_DOCSTRING."""
with ops.name_scope(name, 'RaggedSegmentMean',
[data, segment_ids, num_segments]):
total = segment_sum(data, segment_ids, num_segments)
- ones = ragged_factory_ops.from_nested_row_splits(
- array_ops.ones_like(data.inner_values), data.nested_row_splits)
+ ones = ragged_tensor.RaggedTensor.from_nested_row_splits(
+ array_ops.ones_like(data.flat_values), data.nested_row_splits)
count = segment_sum(ones, segment_ids, num_segments)
- return ragged_factory_ops.from_nested_row_splits(
- total.inner_values / count.inner_values, total.nested_row_splits)
+ if ragged_tensor.is_ragged(total):
+ return total.with_flat_values(total.flat_values / count.flat_values)
+ else:
+ return total / count
def segment_sqrt_n(data, segment_ids, num_segments, name=None):
- # For docs, see: _RAGGED_SEGMENT_DOCSTRING
+ """For docs, see: _RAGGED_SEGMENT_DOCSTRING."""
with ops.name_scope(name, 'RaggedSegmentSqrtN',
[data, segment_ids, num_segments]):
total = segment_sum(data, segment_ids, num_segments)
- ones = ragged_factory_ops.from_nested_row_splits(
- array_ops.ones_like(data.inner_values), data.nested_row_splits)
+ ones = ragged_tensor.RaggedTensor.from_nested_row_splits(
+ array_ops.ones_like(data.flat_values), data.nested_row_splits)
count = segment_sum(ones, segment_ids, num_segments)
- return ragged_factory_ops.from_nested_row_splits(
- total.inner_values / math_ops.sqrt(count.inner_values),
- total.nested_row_splits)
+ if ragged_tensor.is_ragged(total):
+ return total.with_flat_values(
+ total.flat_values / math_ops.sqrt(count.flat_values))
+ else:
+ return total / math_ops.sqrt(count)
def _set_ragged_segment_docstring(func, combination, combined):
@@ -311,7 +320,7 @@
_RAGGED_REDUCE_DOCSTRING = """\
Computes the %(combination)s of elements across dimensions of a `RaggedTensor`.
- Reduces `rt_input` along the dimensions given in `axis` by taking the
+ Reduces `input_tensor` along the dimensions given in `axis` by taking the
%(combination)s of values. If a reduced dimension has no elements for
some index, then the value for that index will be %(default)s.
@@ -319,18 +328,18 @@
`axis` is not specified, then all dimensions are reduced, and a scalar
value is returned.
Args:
- rt_input: A `RaggedTensor` containing the values to be %(combined)s.
+ input_tensor: A `RaggedTensor` containing the values to be %(combined)s.
axis: The dimensions to reduce. May be `None` (to reduce all axes), an
`int` (to reduce a single axis), a `list` or `tuple` of `int` (to reduce
a given set of axes), or a `Tensor` with a constant value. Must be in
- the range `[0, rt_input.rank]`.
+ the range `[0, input_tensor.rank]`.
name: A name prefix for the returned tensor (optional).
Returns:
A `RaggedTensor` containing the %(combined)s values. The returned tensor
has the same dtype as `data`, and its shape is given by removing the
- dimensions specified in `axis` from `rt_input.shape`. The `ragged_rank`
+ dimensions specified in `axis` from `input_tensor.shape`. The `ragged_rank`
of the returned tensor is given by substracting any ragged dimensions
- specified in `axis` from `rt_input.ragged_rank`.
+ specified in `axis` from `input_tensor.ragged_rank`.
Raises:
ValueError: If `axis` contains a `Tensor` whose value is not constant.
####Example:
@@ -387,7 +396,11 @@
"""
-def _ragged_reduce_aggregate(reduce_op, unsorted_segment_op, rt_input, axis,
+def _ragged_reduce_aggregate(reduce_op,
+ unsorted_segment_op,
+ rt_input,
+ axis,
+ keepdims,
name=None):
"""Aggregates across axes of a RaggedTensor using the given `Tensor` ops.
@@ -412,6 +425,7 @@
`int` (to reduce a single axis), a `list` or `tuple` of `int` (to reduce a
given set of axes), or a `Tensor` with a constant value. Must be in the
range `[0, rt_input.rank)`.
+ keepdims: If true, retains reduced dimensions with length 1.
name: A name prefix for the returned tensor (optional).
Returns:
@@ -426,14 +440,19 @@
if not ragged_tensor.is_ragged(rt_input):
return reduce_op(rt_input, axis, name=name)
+ if keepdims:
+ raise ValueError('keepdims=True is not supported for RaggedTensors.')
+
if isinstance(axis, ops.Tensor):
axis = tensor_util.constant_value(axis)
if axis is None:
raise ValueError('axis must be known at graph construction time.')
+ if isinstance(axis, np.ndarray):
+ axis = axis.tolist()
# When reducing all axes, just ignore splits & reduce the inner values.
if axis is None:
- return reduce_op(rt_input.inner_values, None, name=name)
+ return reduce_op(rt_input.flat_values, None, name=name)
with ops.name_scope(name, 'RaggedReduce', [rt_input, axis]):
if isinstance(axis, (tuple, list)):
@@ -448,15 +467,15 @@
# once will probably require a nontrivial c++ op.
axis = sorted(axis)
inner_reduced = _ragged_reduce_aggregate(reduce_op, unsorted_segment_op,
- rt_input, axis[-1])
+ rt_input, axis[-1], keepdims)
return _ragged_reduce_aggregate(reduce_op, unsorted_segment_op,
- inner_reduced, axis[:-1])
+ inner_reduced, axis[:-1], keepdims)
+
+ rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+ rt_input, name='rt_input')
axis = ragged_util.get_positive_axis(axis, rt_input.shape.ndims)
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(
- rt_input, name='rt_input')
-
if axis == 0:
# out[i_1, i_2, ..., i_N] = sum_{j} rt_input[j, i_1, i_2, ..., i_N]
row_lengths = rt_input.row_splits[1:] - rt_input.row_splits[:-1]
@@ -476,69 +495,74 @@
# sum_{j} rt_input [i_0, ..., i_[axis-1], j, i_axis+1], ..., i_N]
return rt_input.with_values(
_ragged_reduce_aggregate(reduce_op, unsorted_segment_op,
- rt_input.values, axis - 1))
+ rt_input.values, axis - 1, keepdims))
-def reduce_sum(rt_input, axis=None, name=None):
+def reduce_sum(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
return _ragged_reduce_aggregate(math_ops.reduce_sum,
- math_ops.unsorted_segment_sum, rt_input, axis,
- name or 'RaggedReduceSum')
+ math_ops.unsorted_segment_sum, input_tensor,
+ axis, keepdims, name or 'RaggedReduceSum')
-def reduce_prod(rt_input, axis=None, name=None):
+def reduce_prod(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
return _ragged_reduce_aggregate(math_ops.reduce_prod,
- math_ops.unsorted_segment_prod, rt_input,
- axis, name or 'RaggedReduceProd')
+ math_ops.unsorted_segment_prod, input_tensor,
+ axis, keepdims, name or 'RaggedReduceProd')
-def reduce_min(rt_input, axis=None, name=None):
+def reduce_min(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
return _ragged_reduce_aggregate(math_ops.reduce_min,
- math_ops.unsorted_segment_min, rt_input, axis,
- name or 'RaggedReduceMin')
+ math_ops.unsorted_segment_min, input_tensor,
+ axis, keepdims, name or 'RaggedReduceMin')
-def reduce_max(rt_input, axis=None, name=None):
+def reduce_max(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
return _ragged_reduce_aggregate(math_ops.reduce_max,
- math_ops.unsorted_segment_max, rt_input, axis,
- name or 'RaggedReduceMax')
+ math_ops.unsorted_segment_max, input_tensor,
+ axis, keepdims, name or 'RaggedReduceMax')
-def reduce_mean(rt_input, axis=None, name=None):
+def reduce_mean(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
- with ops.name_scope(name, 'RaggedReduceMean', [rt_input, axis]):
- total = reduce_sum(rt_input, axis)
- if ragged_tensor.is_ragged(rt_input):
- ones = ragged_factory_ops.from_nested_row_splits(
- array_ops.ones_like(rt_input.inner_values),
- rt_input.nested_row_splits)
+ with ops.name_scope(name, 'RaggedReduceMean', [input_tensor, axis]):
+ total = reduce_sum(input_tensor, axis, keepdims)
+ if ragged_tensor.is_ragged(input_tensor):
+ ones = ragged_tensor.RaggedTensor.from_nested_row_splits(
+ array_ops.ones_like(input_tensor.flat_values),
+ input_tensor.nested_row_splits)
else:
- ones = array_ops.ones_like(rt_input)
- count = reduce_sum(ones, axis)
+ ones = array_ops.ones_like(input_tensor)
+ count = reduce_sum(ones, axis, keepdims)
if ragged_tensor.is_ragged(total):
- return ragged_factory_ops.from_nested_row_splits(
- total.inner_values / count.inner_values, total.nested_row_splits)
+ return ragged_tensor.RaggedTensor.from_nested_row_splits(
+ total.flat_values / count.flat_values, total.nested_row_splits)
else:
return total / count
-def _cast(rt_input, dtype):
- return ragged_functional_ops.map_inner_values(math_ops.cast, rt_input, dtype)
+def _cast(input_tensor, dtype):
+ return ragged_functional_ops.map_flat_values(math_ops.cast, input_tensor,
+ dtype)
-def reduce_all(rt_input, axis=None, name=None):
+def reduce_all(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
- with ops.name_scope(name, 'RaggedReduceAll', [rt_input, axis]):
- return _cast(reduce_prod(_cast(rt_input, dtypes.int32), axis), dtypes.bool)
+ with ops.name_scope(name, 'RaggedReduceAll', [input_tensor, axis]):
+ return _cast(
+ reduce_prod(_cast(input_tensor, dtypes.int32), axis, keepdims),
+ dtypes.bool)
-def reduce_any(rt_input, axis=None, name=None):
+def reduce_any(input_tensor, axis=None, keepdims=None, name=None):
"""For docs, see: _RAGGED_REDUCE_DOCSTRING."""
- with ops.name_scope(name, 'RaggedReduceAny', [rt_input, axis]):
- return _cast(reduce_sum(_cast(rt_input, dtypes.int32), axis), dtypes.bool)
+ with ops.name_scope(name, 'RaggedReduceAny', [input_tensor, axis]):
+ return _cast(
+ reduce_sum(_cast(input_tensor, dtypes.int32), axis, keepdims),
+ dtypes.bool)
def _set_ragged_reduce_docstring(func, combination, combined, default, example):
@@ -554,9 +578,11 @@
_set_ragged_reduce_docstring(reduce_prod, 'product', 'multiplied', '1',
_RAGGED_REDUCE_PROD_EXAMPLE)
_set_ragged_reduce_docstring(reduce_min, 'minimum', 'minimized',
- '`rt_input.dtype.min`', _RAGGED_REDUCE_MIN_EXAMPLE)
+ '`input_tensor.dtype.min`',
+ _RAGGED_REDUCE_MIN_EXAMPLE)
_set_ragged_reduce_docstring(reduce_max, 'maximum', 'maximized',
- '`rt_input.dtype.max`', _RAGGED_REDUCE_MAX_EXAMPLE)
+ '`input_tensor.dtype.max`',
+ _RAGGED_REDUCE_MAX_EXAMPLE)
_set_ragged_reduce_docstring(reduce_mean, 'mean', 'averaged', 'NaN',
_RAGGED_REDUCE_MEAN_EXAMPLE)
diff --git a/tensorflow/python/ops/ragged/ragged_operators.py b/tensorflow/python/ops/ragged/ragged_operators.py
index 223ba0d..7654fa2 100644
--- a/tensorflow/python/ops/ragged/ragged_operators.py
+++ b/tensorflow/python/ops/ragged/ragged_operators.py
@@ -18,7 +18,7 @@
from __future__ import division
from __future__ import print_function
-from tensorflow.python.ops.ragged import ragged_elementwise_ops
+from tensorflow.python.ops import math_ops
from tensorflow.python.ops.ragged import ragged_getitem
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.util import tf_decorator
@@ -33,40 +33,39 @@
ragged_tensor.RaggedTensor.__getitem__ = ragged_getitem.ragged_tensor_getitem
# Ordering operators
-ragged_tensor.RaggedTensor.__ge__ = ragged_elementwise_ops.greater_equal
-ragged_tensor.RaggedTensor.__gt__ = ragged_elementwise_ops.greater
-ragged_tensor.RaggedTensor.__le__ = ragged_elementwise_ops.less_equal
-ragged_tensor.RaggedTensor.__lt__ = ragged_elementwise_ops.less
+ragged_tensor.RaggedTensor.__ge__ = math_ops.greater_equal
+ragged_tensor.RaggedTensor.__gt__ = math_ops.greater
+ragged_tensor.RaggedTensor.__le__ = math_ops.less_equal
+ragged_tensor.RaggedTensor.__lt__ = math_ops.less
# Logical operators
-ragged_tensor.RaggedTensor.__and__ = ragged_elementwise_ops.logical_and
-ragged_tensor.RaggedTensor.__rand__ = _right(ragged_elementwise_ops.logical_and)
-ragged_tensor.RaggedTensor.__invert__ = ragged_elementwise_ops.logical_not
-ragged_tensor.RaggedTensor.__ror__ = _right(ragged_elementwise_ops.logical_or)
-ragged_tensor.RaggedTensor.__or__ = ragged_elementwise_ops.logical_or
-ragged_tensor.RaggedTensor.__xor__ = ragged_elementwise_ops.logical_xor
-ragged_tensor.RaggedTensor.__rxor__ = _right(ragged_elementwise_ops.logical_xor)
+ragged_tensor.RaggedTensor.__and__ = math_ops.logical_and
+ragged_tensor.RaggedTensor.__rand__ = _right(math_ops.logical_and)
+ragged_tensor.RaggedTensor.__invert__ = math_ops.logical_not
+ragged_tensor.RaggedTensor.__ror__ = _right(math_ops.logical_or)
+ragged_tensor.RaggedTensor.__or__ = math_ops.logical_or
+ragged_tensor.RaggedTensor.__xor__ = math_ops.logical_xor
+ragged_tensor.RaggedTensor.__rxor__ = _right(math_ops.logical_xor)
# Arithmetic operators
-ragged_tensor.RaggedTensor.__abs__ = ragged_elementwise_ops.abs
-ragged_tensor.RaggedTensor.__add__ = ragged_elementwise_ops.add
-ragged_tensor.RaggedTensor.__radd__ = _right(ragged_elementwise_ops.add)
-ragged_tensor.RaggedTensor.__div__ = ragged_elementwise_ops.div
-ragged_tensor.RaggedTensor.__rdiv__ = _right(ragged_elementwise_ops.div)
-ragged_tensor.RaggedTensor.__floordiv__ = ragged_elementwise_ops.floordiv
-ragged_tensor.RaggedTensor.__rfloordiv__ = _right(
- ragged_elementwise_ops.floordiv)
-ragged_tensor.RaggedTensor.__mod__ = ragged_elementwise_ops.floormod
-ragged_tensor.RaggedTensor.__rmod__ = _right(ragged_elementwise_ops.floormod)
-ragged_tensor.RaggedTensor.__mul__ = ragged_elementwise_ops.multiply
-ragged_tensor.RaggedTensor.__rmul__ = _right(ragged_elementwise_ops.multiply)
-ragged_tensor.RaggedTensor.__neg__ = ragged_elementwise_ops.negative
-ragged_tensor.RaggedTensor.__pow__ = ragged_elementwise_ops.pow
-ragged_tensor.RaggedTensor.__rpow__ = _right(ragged_elementwise_ops.pow)
-ragged_tensor.RaggedTensor.__sub__ = ragged_elementwise_ops.subtract
-ragged_tensor.RaggedTensor.__rsub__ = _right(ragged_elementwise_ops.subtract)
-ragged_tensor.RaggedTensor.__truediv__ = ragged_elementwise_ops.truediv
-ragged_tensor.RaggedTensor.__rtruediv__ = _right(ragged_elementwise_ops.truediv)
+ragged_tensor.RaggedTensor.__abs__ = math_ops.abs
+ragged_tensor.RaggedTensor.__add__ = math_ops.add
+ragged_tensor.RaggedTensor.__radd__ = _right(math_ops.add)
+ragged_tensor.RaggedTensor.__div__ = math_ops.div
+ragged_tensor.RaggedTensor.__rdiv__ = _right(math_ops.div)
+ragged_tensor.RaggedTensor.__floordiv__ = math_ops.floordiv
+ragged_tensor.RaggedTensor.__rfloordiv__ = _right(math_ops.floordiv)
+ragged_tensor.RaggedTensor.__mod__ = math_ops.floormod
+ragged_tensor.RaggedTensor.__rmod__ = _right(math_ops.floormod)
+ragged_tensor.RaggedTensor.__mul__ = math_ops.multiply
+ragged_tensor.RaggedTensor.__rmul__ = _right(math_ops.multiply)
+ragged_tensor.RaggedTensor.__neg__ = math_ops.negative
+ragged_tensor.RaggedTensor.__pow__ = math_ops.pow
+ragged_tensor.RaggedTensor.__rpow__ = _right(math_ops.pow)
+ragged_tensor.RaggedTensor.__sub__ = math_ops.subtract
+ragged_tensor.RaggedTensor.__rsub__ = _right(math_ops.subtract)
+ragged_tensor.RaggedTensor.__truediv__ = math_ops.truediv
+ragged_tensor.RaggedTensor.__rtruediv__ = _right(math_ops.truediv)
# Dummy methods
diff --git a/tensorflow/python/ops/ragged/ragged_operators_test.py b/tensorflow/python/ops/ragged/ragged_operators_test.py
index 7fe8159..d1c6b90 100644
--- a/tensorflow/python/ops/ragged/ragged_operators_test.py
+++ b/tensorflow/python/ops/ragged/ragged_operators_test.py
@@ -19,85 +19,75 @@
from __future__ import print_function
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase):
- # @TODO(edloper): Test right-handed versions of operators once we add
- # broadcasting support for elementwise ops.
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedElementwiseOpsTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testOrderingOperators(self):
- x = ragged.constant([[1, 5], [3]])
- y = ragged.constant([[4, 5], [1]])
- with self.test_session():
- self.assertEqual((x > y).eval().tolist(), [[False, False], [True]])
- self.assertEqual((x >= y).eval().tolist(), [[False, True], [True]])
- self.assertEqual((x < y).eval().tolist(), [[True, False], [False]])
- self.assertEqual((x <= y).eval().tolist(), [[True, True], [False]])
+ x = ragged_factory_ops.constant([[1, 5], [3]])
+ y = ragged_factory_ops.constant([[4, 5], [1]])
+ self.assertRaggedEqual((x > y), [[False, False], [True]])
+ self.assertRaggedEqual((x >= y), [[False, True], [True]])
+ self.assertRaggedEqual((x < y), [[True, False], [False]])
+ self.assertRaggedEqual((x <= y), [[True, True], [False]])
- def assertEqual(self, a, b):
- if a != b:
- print('%30s %s' % (b, a))
-
- @test_util.run_deprecated_v1
def testArithmeticOperators(self):
- x = ragged.constant([[1.0, -2.0], [8.0]])
- y = ragged.constant([[4.0, 4.0], [2.0]])
- with self.test_session():
- self.assertEqual(abs(x).eval().tolist(), [[1.0, 2.0], [8.0]])
+ x = ragged_factory_ops.constant([[1.0, -2.0], [8.0]])
+ y = ragged_factory_ops.constant([[4.0, 4.0], [2.0]])
+ self.assertRaggedEqual(abs(x), [[1.0, 2.0], [8.0]])
- self.assertEqual((-x).eval().tolist(), [[-1.0, 2.0], [-8.0]])
+ self.assertRaggedEqual((-x), [[-1.0, 2.0], [-8.0]])
- self.assertEqual((x + y).eval().tolist(), [[5.0, 2.0], [10.0]])
- self.assertEqual((3.0 + y).eval().tolist(), [[7.0, 7.0], [5.0]])
- self.assertEqual((x + 3.0).eval().tolist(), [[4.0, 1.0], [11.0]])
+ self.assertRaggedEqual((x + y), [[5.0, 2.0], [10.0]])
+ self.assertRaggedEqual((3.0 + y), [[7.0, 7.0], [5.0]])
+ self.assertRaggedEqual((x + 3.0), [[4.0, 1.0], [11.0]])
- self.assertEqual((x - y).eval().tolist(), [[-3.0, -6.0], [6.0]])
- self.assertEqual((3.0 - y).eval().tolist(), [[-1.0, -1.0], [1.0]])
- self.assertEqual((x + 3.0).eval().tolist(), [[4.0, 1.0], [11.0]])
+ self.assertRaggedEqual((x - y), [[-3.0, -6.0], [6.0]])
+ self.assertRaggedEqual((3.0 - y), [[-1.0, -1.0], [1.0]])
+ self.assertRaggedEqual((x + 3.0), [[4.0, 1.0], [11.0]])
- self.assertEqual((x * y).eval().tolist(), [[4.0, -8.0], [16.0]])
- self.assertEqual((3.0 * y).eval().tolist(), [[12.0, 12.0], [6.0]])
- self.assertEqual((x * 3.0).eval().tolist(), [[3.0, -6.0], [24.0]])
+ self.assertRaggedEqual((x * y), [[4.0, -8.0], [16.0]])
+ self.assertRaggedEqual((3.0 * y), [[12.0, 12.0], [6.0]])
+ self.assertRaggedEqual((x * 3.0), [[3.0, -6.0], [24.0]])
- self.assertEqual((x / y).eval().tolist(), [[0.25, -0.5], [4.0]])
- self.assertEqual((y / x).eval().tolist(), [[4.0, -2.0], [0.25]])
- self.assertEqual((2.0 / y).eval().tolist(), [[0.5, 0.5], [1.0]])
- self.assertEqual((x / 2.0).eval().tolist(), [[0.5, -1.0], [4.0]])
+ self.assertRaggedEqual((x / y), [[0.25, -0.5], [4.0]])
+ self.assertRaggedEqual((y / x), [[4.0, -2.0], [0.25]])
+ self.assertRaggedEqual((2.0 / y), [[0.5, 0.5], [1.0]])
+ self.assertRaggedEqual((x / 2.0), [[0.5, -1.0], [4.0]])
- self.assertEqual((x // y).eval().tolist(), [[0.0, -1.0], [4.0]])
- self.assertEqual((y // x).eval().tolist(), [[4.0, -2.0], [0.0]])
- self.assertEqual((2.0 // y).eval().tolist(), [[0.0, 0.0], [1.0]])
- self.assertEqual((x // 2.0).eval().tolist(), [[0.0, -1.0], [4.0]])
+ self.assertRaggedEqual((x // y), [[0.0, -1.0], [4.0]])
+ self.assertRaggedEqual((y // x), [[4.0, -2.0], [0.0]])
+ self.assertRaggedEqual((2.0 // y), [[0.0, 0.0], [1.0]])
+ self.assertRaggedEqual((x // 2.0), [[0.0, -1.0], [4.0]])
- self.assertEqual((x % y).eval().tolist(), [[1.0, 2.0], [0.0]])
- self.assertEqual((y % x).eval().tolist(), [[0.0, -0.0], [2.0]])
- self.assertEqual((2.0 % y).eval().tolist(), [[2.0, 2.0], [0.0]])
- self.assertEqual((x % 2.0).eval().tolist(), [[1.0, 0.0], [0.0]])
+ self.assertRaggedEqual((x % y), [[1.0, 2.0], [0.0]])
+ self.assertRaggedEqual((y % x), [[0.0, -0.0], [2.0]])
+ self.assertRaggedEqual((2.0 % y), [[2.0, 2.0], [0.0]])
+ self.assertRaggedEqual((x % 2.0), [[1.0, 0.0], [0.0]])
- @test_util.run_deprecated_v1
def testLogicalOperators(self):
- a = ragged.constant([[True, True], [False]])
- b = ragged.constant([[True, False], [False]])
- with self.test_session():
- self.assertEqual((~a).eval().tolist(), [[False, False], [True]])
+ a = ragged_factory_ops.constant([[True, True], [False]])
+ b = ragged_factory_ops.constant([[True, False], [False]])
+ self.assertRaggedEqual((~a), [[False, False], [True]])
- self.assertEqual((a & b).eval().tolist(), [[True, False], [False]])
- self.assertEqual((a & True).eval().tolist(), [[True, True], [False]])
- self.assertEqual((True & b).eval().tolist(), [[True, False], [False]])
+ self.assertRaggedEqual((a & b), [[True, False], [False]])
+ self.assertRaggedEqual((a & True), [[True, True], [False]])
+ self.assertRaggedEqual((True & b), [[True, False], [False]])
- self.assertEqual((a | b).eval().tolist(), [[True, True], [False]])
- self.assertEqual((a | False).eval().tolist(), [[True, True], [False]])
- self.assertEqual((False | b).eval().tolist(), [[True, False], [False]])
+ self.assertRaggedEqual((a | b), [[True, True], [False]])
+ self.assertRaggedEqual((a | False), [[True, True], [False]])
+ self.assertRaggedEqual((False | b), [[True, False], [False]])
- self.assertEqual((a ^ b).eval().tolist(), [[False, True], [False]])
- self.assertEqual((a ^ True).eval().tolist(), [[False, False], [True]])
- self.assertEqual((True ^ b).eval().tolist(), [[False, True], [True]])
+ self.assertRaggedEqual((a ^ b), [[False, True], [False]])
+ self.assertRaggedEqual((a ^ True), [[False, False], [True]])
+ self.assertRaggedEqual((True ^ b), [[False, True], [True]])
def testDummyOperators(self):
- a = ragged.constant([[True, True], [False]])
+ a = ragged_factory_ops.constant([[True, True], [False]])
with self.assertRaisesRegexp(TypeError,
'RaggedTensor may not be used as a boolean.'):
bool(a)
diff --git a/tensorflow/python/ops/ragged/ragged_range_op_test.py b/tensorflow/python/ops/ragged/ragged_range_op_test.py
index 644423e..afe5866 100644
--- a/tensorflow/python/ops/ragged/ragged_range_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_range_op_test.py
@@ -20,113 +20,108 @@
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedRangeOpTest(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedRangeOpTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testDocStringExamples(self):
"""Examples from ragged_range.__doc__."""
- with self.test_session():
- rt1 = ragged.range([3, 5, 2]).eval().tolist()
- self.assertEqual(rt1, [[0, 1, 2], [0, 1, 2, 3, 4], [0, 1]])
+ rt1 = ragged_math_ops.range([3, 5, 2])
+ self.assertRaggedEqual(rt1, [[0, 1, 2], [0, 1, 2, 3, 4], [0, 1]])
- rt2 = ragged.range([0, 5, 8], [3, 3, 12]).eval().tolist()
- self.assertEqual(rt2, [[0, 1, 2], [], [8, 9, 10, 11]])
+ rt2 = ragged_math_ops.range([0, 5, 8], [3, 3, 12])
+ self.assertRaggedEqual(rt2, [[0, 1, 2], [], [8, 9, 10, 11]])
- rt3 = ragged.range([0, 5, 8], [3, 3, 12], 2).eval().tolist()
- self.assertEqual(rt3, [[0, 2], [], [8, 10]])
+ rt3 = ragged_math_ops.range([0, 5, 8], [3, 3, 12], 2)
+ self.assertRaggedEqual(rt3, [[0, 2], [], [8, 10]])
- @test_util.run_deprecated_v1
def testBasicRanges(self):
- with self.test_session():
- # Specify limits only.
- self.assertEqual(
- ragged.range([0, 3, 5]).eval().tolist(),
- [list(range(0)), list(range(3)), list(range(5))])
+ # Specify limits only.
+ self.assertRaggedEqual(
+ ragged_math_ops.range([0, 3, 5]),
+ [list(range(0)), list(range(3)),
+ list(range(5))])
- # Specify starts and limits.
- self.assertEqual(
- ragged.range([0, 3, 5], [2, 3, 10]).eval().tolist(),
- [list(range(0, 2)), list(range(3, 3)), list(range(5, 10))])
+ # Specify starts and limits.
+ self.assertRaggedEqual(
+ ragged_math_ops.range([0, 3, 5], [2, 3, 10]),
+ [list(range(0, 2)),
+ list(range(3, 3)),
+ list(range(5, 10))])
- # Specify starts, limits, and deltas.
- self.assertEqual(
- ragged.range([0, 3, 5], [4, 4, 15], [2, 3, 4]).eval().tolist(),
- [list(range(0, 4, 2)), list(range(3, 4, 3)),
- list(range(5, 15, 4))])
+ # Specify starts, limits, and deltas.
+ self.assertRaggedEqual(
+ ragged_math_ops.range([0, 3, 5], [4, 4, 15], [2, 3, 4]),
+ [list(range(0, 4, 2)),
+ list(range(3, 4, 3)),
+ list(range(5, 15, 4))])
- @test_util.run_deprecated_v1
def testFloatRanges(self):
- with self.test_session():
- expected = [[0.0, 0.4, 0.8, 1.2, 1.6, 2.0, 2.4, 2.8, 3.2, 3.6], [3.0],
- [5.0, 7.2, 9.4, 11.6, 13.8]]
- actual = ragged.range([0.0, 3.0, 5.0], [3.9, 4.0, 15.0],
- [0.4, 1.5, 2.2]).eval().tolist()
- self.assertEqual(expected, [[round(v, 5) for v in row] for row in actual])
-
- @test_util.run_deprecated_v1
- def testNegativeDeltas(self):
- with self.test_session():
- self.assertEqual(
- ragged.range([0, 3, 5], limits=0, deltas=-1).eval().tolist(),
- [list(range(0, 0, -1)), list(range(3, 0, -1)),
- list(range(5, 0, -1))])
-
- self.assertEqual(
- ragged.range([0, -3, 5], limits=0, deltas=[-1, 1,
- -2]).eval().tolist(),
- [list(range(0, 0, -1)), list(range(-3, 0, 1)),
- list(range(5, 0, -2))])
-
- @test_util.run_deprecated_v1
- def testBroadcast(self):
- with self.test_session():
- # Specify starts and limits, broadcast deltas.
- self.assertEqual(
- ragged.range([0, 3, 5], [4, 4, 15], 3).eval().tolist(),
- [list(range(0, 4, 3)), list(range(3, 4, 3)),
- list(range(5, 15, 3))])
-
- # Broadcast all arguments.
- self.assertEqual(
- ragged.range(0, 5, 1).eval().tolist(), [list(range(0, 5, 1))])
-
- @test_util.run_deprecated_v1
- def testEmptyRanges(self):
- rt1 = ragged.range([0, 5, 3], [0, 3, 5])
- rt2 = ragged.range([0, 5, 5], [0, 3, 5], -1)
- with self.test_session():
- self.assertEqual(rt1.eval().tolist(), [[], [], [3, 4]])
- self.assertEqual(rt2.eval().tolist(), [[], [5, 4], []])
-
- @test_util.run_deprecated_v1
- def testShapeFnErrors(self):
- with self.test_session():
- self.assertRaisesRegexp(ValueError, r'Shape must be at most rank 1.*',
- ragged.range, [[0]], 5)
- self.assertRaisesRegexp(ValueError, r'Shape must be at most rank 1.*',
- ragged.range, 0, [[5]])
- self.assertRaisesRegexp(ValueError, r'Shape must be at most rank 1.*',
- ragged.range, 0, 5, [[0]])
- self.assertRaisesRegexp(ValueError, r'Dimensions must be equal.*',
- ragged.range, [0], [1, 2])
-
- @test_util.run_deprecated_v1
- def testKernelErrors(self):
- with self.test_session():
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- r'Requires delta != 0',
- ragged.range(0, 0, 0).eval)
-
- @test_util.run_deprecated_v1
- def testShape(self):
- self.assertEqual(ragged.range(0, 0, 0).shape.as_list(), [1, None])
- self.assertEqual(ragged.range([1, 2, 3]).shape.as_list(), [3, None])
+ expected = [[0.0, 0.4, 0.8, 1.2, 1.6, 2.0, 2.4, 2.8, 3.2, 3.6], [3.0],
+ [5.0, 7.2, 9.4, 11.6, 13.8]]
+ actual = ragged_math_ops.range([0.0, 3.0, 5.0], [3.9, 4.0, 15.0],
+ [0.4, 1.5, 2.2])
self.assertEqual(
- ragged.range([1, 2, 3], [4, 5, 6]).shape.as_list(), [3, None])
+ expected,
+ [[round(v, 5) for v in row] for row in self.eval_to_list(actual)])
+
+ def testNegativeDeltas(self):
+ self.assertRaggedEqual(
+ ragged_math_ops.range([0, 3, 5], limits=0, deltas=-1),
+ [list(range(0, 0, -1)),
+ list(range(3, 0, -1)),
+ list(range(5, 0, -1))])
+
+ self.assertRaggedEqual(
+ ragged_math_ops.range([0, -3, 5], limits=0, deltas=[-1, 1, -2]),
+ [list(range(0, 0, -1)),
+ list(range(-3, 0, 1)),
+ list(range(5, 0, -2))])
+
+ def testBroadcast(self):
+ # Specify starts and limits, broadcast deltas.
+ self.assertRaggedEqual(
+ ragged_math_ops.range([0, 3, 5], [4, 4, 15], 3),
+ [list(range(0, 4, 3)),
+ list(range(3, 4, 3)),
+ list(range(5, 15, 3))])
+
+ # Broadcast all arguments.
+ self.assertRaggedEqual(
+ ragged_math_ops.range(0, 5, 1), [list(range(0, 5, 1))])
+
+ def testEmptyRanges(self):
+ rt1 = ragged_math_ops.range([0, 5, 3], [0, 3, 5])
+ rt2 = ragged_math_ops.range([0, 5, 5], [0, 3, 5], -1)
+ self.assertRaggedEqual(rt1, [[], [], [3, 4]])
+ self.assertRaggedEqual(rt2, [[], [5, 4], []])
+
+ def testShapeFnErrors(self):
+ self.assertRaises((ValueError, errors.InvalidArgumentError),
+ ragged_math_ops.range, [[0]], 5)
+ self.assertRaises((ValueError, errors.InvalidArgumentError),
+ ragged_math_ops.range, 0, [[5]])
+ self.assertRaises((ValueError, errors.InvalidArgumentError),
+ ragged_math_ops.range, 0, 5, [[0]])
+ self.assertRaises((ValueError, errors.InvalidArgumentError),
+ ragged_math_ops.range, [0], [1, 2])
+
+ def testKernelErrors(self):
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ r'Requires delta != 0'):
+ self.evaluate(ragged_math_ops.range(0, 0, 0))
+
+ def testShape(self):
+ self.assertRaggedEqual(
+ ragged_math_ops.range(0, 0, 1).shape.as_list(), [1, None])
+ self.assertRaggedEqual(
+ ragged_math_ops.range([1, 2, 3]).shape.as_list(), [3, None])
+ self.assertRaggedEqual(
+ ragged_math_ops.range([1, 2, 3], [4, 5, 6]).shape.as_list(), [3, None])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_reduce_op_test.py b/tensorflow/python/ops/ragged/ragged_reduce_op_test.py
index 9f51d59..a9fa378 100644
--- a/tensorflow/python/ops/ragged/ragged_reduce_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_reduce_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.reduce_<AGGREGATE> ops."""
+"""Tests for ragged_math_ops.reduce_<AGGREGATE> ops."""
from __future__ import absolute_import
from __future__ import division
@@ -21,11 +21,14 @@
from absl.testing import parameterized
import numpy as np
+from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
_MAX_INT32 = dtypes.int32.max
@@ -37,7 +40,9 @@
return 1.0 * sum(values) / len(values)
-class RaggedReduceOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedReduceOpsTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters(
#=========================================================================
@@ -48,88 +53,88 @@
# [2, 6 ]]
#=========================================================================
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=0,
expected=[15, 12, 4] # = [3+1+9+2, 1+5+6, 4]
),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=-2,
expected=[15, 12, 4] # = [3+1+9+2, 1+5+6, 4]
),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=1,
expected=[8, 6, 9, 8] # = [3+1+4, 1+5, 9, 2+6]
),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=-1,
expected=[8, 6, 9, 8] # = [3+1+4, 1+5, 9, 2+6]
),
dict(
- ragged_reduce_op=ragged.reduce_prod,
+ ragged_reduce_op=ragged_math_ops.reduce_prod,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=0,
expected=[54, 30, 4] # = [3*1*9*2, 1*5*6, 4]
),
dict(
- ragged_reduce_op=ragged.reduce_prod,
+ ragged_reduce_op=ragged_math_ops.reduce_prod,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=1,
expected=[12, 5, 9, 12] # = [3*1*4, 1*5, 9, 2*6]
),
dict(
- ragged_reduce_op=ragged.reduce_min,
+ ragged_reduce_op=ragged_math_ops.reduce_min,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=0,
expected=[1, 1, 4] # = [min(3, 1, 9, 2), min(1, 5, 6), 4]
),
dict(
- ragged_reduce_op=ragged.reduce_min,
+ ragged_reduce_op=ragged_math_ops.reduce_min,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=1,
expected=[1, 1, 9, 2] # = [min(3, 1, 4), min(1, 5), 9, min(2, 6)]
),
dict(
- ragged_reduce_op=ragged.reduce_max,
+ ragged_reduce_op=ragged_math_ops.reduce_max,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=0,
expected=[9, 6, 4] # = [max(3, 1, 9, 2), max(1, 5, 6), 4]
),
dict(
- ragged_reduce_op=ragged.reduce_max,
+ ragged_reduce_op=ragged_math_ops.reduce_max,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=1,
expected=[4, 5, 9, 6] # = [max(3, 1, 4), max(1, 5), 9, max(2, 6)]
),
dict(
- ragged_reduce_op=ragged.reduce_mean,
+ ragged_reduce_op=ragged_math_ops.reduce_mean,
rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
axis=0,
expected=[3.75, 4, 4] # = [mean(3, 1, 9, 2), mean(1, 5, 6), 4]
),
dict(
- ragged_reduce_op=ragged.reduce_any,
+ ragged_reduce_op=ragged_math_ops.reduce_any,
rt_input=[[True, True], [True, True, False, True], [False, True]],
axis=0,
expected=[True, True, False, True]),
dict(
- ragged_reduce_op=ragged.reduce_any,
+ ragged_reduce_op=ragged_math_ops.reduce_any,
rt_input=[[True, True], [True, True, False, True], [False, True]],
axis=1,
expected=[True, True, True]),
dict(
- ragged_reduce_op=ragged.reduce_all,
+ ragged_reduce_op=ragged_math_ops.reduce_all,
rt_input=[[True, True], [True, True, False, True], [False, True]],
axis=0,
expected=[False, True, False, True]),
dict(
- ragged_reduce_op=ragged.reduce_all,
+ ragged_reduce_op=ragged_math_ops.reduce_all,
rt_input=[[True, True], [True, True, False, True], [False, True]],
axis=1,
expected=[True, False, False]),
@@ -146,53 +151,53 @@
# axis=None
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=None,
expected=0 + 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9),
dict(
- ragged_reduce_op=ragged.reduce_prod,
+ ragged_reduce_op=ragged_math_ops.reduce_prod,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=None,
expected=0 * 1 * 2 * 3 * 4 * 5 * 6 * 7 * 8 * 9),
dict(
- ragged_reduce_op=ragged.reduce_min,
+ ragged_reduce_op=ragged_math_ops.reduce_min,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=None,
expected=min(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
dict(
- ragged_reduce_op=ragged.reduce_max,
+ ragged_reduce_op=ragged_math_ops.reduce_max,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=None,
expected=max(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
dict(
- ragged_reduce_op=ragged.reduce_mean,
+ ragged_reduce_op=ragged_math_ops.reduce_mean,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=None,
expected=mean(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
# axis=0
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=0,
expected=[0 + 4 + 5 + 7 + 8, 1 + 6 + 9, 2, 3]),
dict(
- ragged_reduce_op=ragged.reduce_prod,
+ ragged_reduce_op=ragged_math_ops.reduce_prod,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=0,
expected=[0 * 4 * 5 * 7 * 8, 1 * 6 * 9, 2, 3]),
dict(
- ragged_reduce_op=ragged.reduce_min,
+ ragged_reduce_op=ragged_math_ops.reduce_min,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=0,
expected=[min(0, 4, 5, 7, 8), min(1, 6, 9), 2, 3]),
dict(
- ragged_reduce_op=ragged.reduce_max,
+ ragged_reduce_op=ragged_math_ops.reduce_max,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=0,
expected=[max(0, 4, 5, 7, 8), max(1, 6, 9), 2, 3]),
dict(
- ragged_reduce_op=ragged.reduce_mean,
+ ragged_reduce_op=ragged_math_ops.reduce_mean,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=0,
expected=[mean(0, 4, 5, 7, 8),
@@ -201,24 +206,24 @@
# Note: we don't test mean here because it gives a NaN, and this will
# cause assertEqual to fail (since NaN != NaN). See testMeanNan().
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=1,
expected=[0 + 1 + 2 + 3, 4, 0, 5 + 6, 7, 8 + 9]),
dict(
- ragged_reduce_op=ragged.reduce_prod,
+ ragged_reduce_op=ragged_math_ops.reduce_prod,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=1,
expected=[0 * 1 * 2 * 3, 4, 1, 5 * 6, 7, 8 * 9]),
dict(
- ragged_reduce_op=ragged.reduce_min,
+ ragged_reduce_op=ragged_math_ops.reduce_min,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=1,
expected=[min(0, 1, 2, 3), 4, _MAX_INT32,
min(5, 6), 7,
min(8, 9)]),
dict(
- ragged_reduce_op=ragged.reduce_max,
+ ragged_reduce_op=ragged_math_ops.reduce_max,
rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
axis=1,
expected=[max(0, 1, 2, 3), 4, _MIN_INT32,
@@ -233,47 +238,47 @@
# [[9 ] ]]
#=========================================================================
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=[],
expected=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=None,
expected=sum([1, 2, 3, 4, 5, 6, 7, 8, 9])),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=0,
expected=[[1 + 6 + 9, 2 + 7], [], [3 + 8, 4, 5]]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=1,
expected=[[1 + 3, 2 + 4, 5], [6 + 8, 7], [], [9]]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=2,
expected=[[1 + 2, 0, 3 + 4 + 5], [6 + 7, 0, 8], [], [9]]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=[0, 1],
expected=[1 + 3 + 6 + 8 + 9, 2 + 4 + 7, 5]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=[0, 2],
expected=[1 + 6 + 9 + 2 + 7, 0, 3 + 8 + 4 + 5]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=[1, 2],
expected=[1 + 2 + 3 + 4 + 5, 6 + 7 + 8, 0, 9]),
dict(
- ragged_reduce_op=ragged.reduce_sum,
+ ragged_reduce_op=ragged_math_ops.reduce_sum,
rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
axis=[0, 1, 2],
expected=sum([1, 2, 3, 4, 5, 6, 7, 8, 9])),
@@ -285,62 +290,56 @@
# [[9 ] ]]
#=========================================================================
dict(
- ragged_reduce_op=ragged.reduce_mean,
+ ragged_reduce_op=ragged_math_ops.reduce_mean,
rt_input=[[[1, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
axis=0,
expected=[[mean(1, 6, 9), mean(2, 7)], [mean(3, 8), 4, 5]]),
dict(
- ragged_reduce_op=ragged.reduce_mean,
+ ragged_reduce_op=ragged_math_ops.reduce_mean,
rt_input=[[[1, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
axis=1,
expected=[[mean(1, 3), mean(2, 4), 5], [mean(6, 8), 7], [9]]),
dict(
- ragged_reduce_op=ragged.reduce_mean,
+ ragged_reduce_op=ragged_math_ops.reduce_mean,
rt_input=[[[1, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
axis=2,
expected=[[mean(1, 2), mean(3, 4, 5)], [mean(6, 7), 8], [9]]),
)
- @test_util.run_deprecated_v1
def testReduce(self, ragged_reduce_op, rt_input, axis, expected):
- rt_input = ragged.constant(rt_input)
+ rt_input = ragged_factory_ops.constant(rt_input)
reduced = ragged_reduce_op(rt_input, axis)
- with self.test_session():
- self.assertEqual(reduced.eval().tolist(), expected)
+ self.assertRaggedEqual(reduced, expected)
def assertEqualWithNan(self, actual, expected):
"""Like assertEqual, but NaN==NaN."""
self.assertTrue(
((actual == expected) | (np.isnan(actual) & np.isnan(expected))).all())
- @test_util.run_deprecated_v1
def testMeanNan(self):
rt_as_list = [[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]]
expected = (
np.array([0 + 1 + 2 + 3, 4, 0, 5 + 6, 7, 8 + 9]) / np.array(
[4, 1, 0, 2, 1, 2]))
- rt_input = ragged.constant(rt_as_list)
- reduced = ragged.reduce_mean(rt_input, axis=1)
- with self.test_session():
- self.assertEqualWithNan(reduced.eval(), expected)
+ rt_input = ragged_factory_ops.constant(rt_as_list)
+ reduced = ragged_math_ops.reduce_mean(rt_input, axis=1)
+ self.assertEqualWithNan(self.evaluate(reduced), expected)
- @test_util.run_deprecated_v1
def testMeanWithTensorInputs(self):
tensor = [[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]
expected = [2.0, 20.0]
- reduced = ragged.reduce_mean(tensor, axis=1)
- with self.test_session():
- self.assertAllEqual(reduced.eval(), expected)
+ reduced = ragged_math_ops.reduce_mean(tensor, axis=1)
+ self.assertRaggedEqual(reduced, expected)
- @test_util.run_deprecated_v1
def testErrors(self):
- rt_input = ragged.constant([[1, 2, 3], [4, 5]])
+ rt_input = ragged_factory_ops.constant([[1, 2, 3], [4, 5]])
axis = array_ops.placeholder_with_default(constant_op.constant([0]), None)
- self.assertRaisesRegexp(ValueError,
- r'axis must be known at graph construction time.',
- ragged.reduce_sum, rt_input, axis)
- self.assertRaisesRegexp(TypeError,
- r'axis must be an int; got str.*',
- ragged.reduce_sum, rt_input, ['x'])
+
+ if not context.executing_eagerly():
+ self.assertRaisesRegexp(
+ ValueError, r'axis must be known at graph construction time.',
+ ragged_math_ops.reduce_sum, rt_input, axis)
+ self.assertRaisesRegexp(TypeError, r'axis must be an int; got str.*',
+ ragged_math_ops.reduce_sum, rt_input, ['x'])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_row_lengths_op_test.py b/tensorflow/python/ops/ragged/ragged_row_lengths_op_test.py
index 4a705be..8f8089c 100644
--- a/tensorflow/python/ops/ragged/ragged_row_lengths_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_row_lengths_op_test.py
@@ -20,12 +20,17 @@
from absl.testing import parameterized
+from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedRowLengthsOp(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedRowLengthsOp(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters([
# Docstring Example
@@ -37,24 +42,6 @@
axis=2,
expected=[[3, 1], [], [2, 1], [1], []]),
- # 1D tensor
- dict(
- rt_input=[1, 2, 3, 4, 5],
- ragged_rank=0,
- axis=0,
- expected=5),
-
- # 2D Tensor (0 ragged dimensions)
- dict(
- rt_input=[[1, 2], [3, 4], [5, 6], [7, 8]],
- ragged_rank=0,
- expected=[2, 2, 2, 2]),
- dict(
- rt_input=[[1, 2], [3, 4], [5, 6], [7, 8]],
- ragged_rank=0,
- axis=0,
- expected=4),
-
# 2D Tensor (1 ragged dimension)
dict(
rt_input=[['a'], ['b', 'c', 'd'], ['e'], [], ['f']],
@@ -79,24 +66,6 @@
axis=0,
expected=0),
- # 3D Tensor (0 ragged dimensions)
- dict(
- rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]],
- ragged_rank=0,
- axis=0,
- expected=2),
- dict(
- rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]],
- ragged_rank=0,
- axis=1,
- expected=[3, 3]),
- dict(
- rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]],
- ragged_rank=0,
- axis=2,
- expected=[[2, 2, 2], [2, 2, 2]],
- expected_ragged_rank=0),
-
# 3D Tensor (1 ragged dimension)
dict(
rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10]]],
@@ -143,42 +112,35 @@
expected=[[2, 3, 0], [4, 1]],
expected_ragged_rank=1),
]) # pyformat: disable
- @test_util.run_deprecated_v1
def testRowLengths(self,
rt_input,
expected,
axis=1,
ragged_rank=None,
expected_ragged_rank=None):
- rt = ragged.constant(rt_input, ragged_rank=ragged_rank)
- lengths = ragged.row_lengths(rt, axis)
- with self.test_session():
- self.assertEqual(lengths.eval().tolist(), expected)
- if expected_ragged_rank is not None:
- if isinstance(lengths, ragged.RaggedTensor):
- self.assertEqual(lengths.ragged_rank, expected_ragged_rank)
- else:
- self.assertEqual(0, expected_ragged_rank)
+ rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
+ lengths = rt.row_lengths(axis)
+ self.assertRaggedEqual(lengths, expected)
+ if expected_ragged_rank is not None:
+ if isinstance(lengths, ragged_tensor.RaggedTensor):
+ self.assertEqual(lengths.ragged_rank, expected_ragged_rank)
+ else:
+ self.assertEqual(0, expected_ragged_rank)
@parameterized.parameters([
- dict(
- rt_input=10,
- exception=ValueError,
- message='rt_input may not be a scalar.'),
- dict(
- rt_input=[10, 20],
- axis=1,
- exception=ValueError,
- message='axis=1 out of bounds: expected -1<=axis<1.'),
- dict(
+ dict( # axis=2 out of bounds: expected -2<=axis<2.
+ rt_input=[[10, 20], [30]],
+ axis=2,
+ exception=(ValueError, errors.InvalidArgumentError)),
+ dict( # axis=-3 out of bounds: expected -2<=axis<2.
rt_input=[[2, 3, 0], [4, 1, 2]],
axis=-3,
- exception=ValueError,
- message='axis=-3 out of bounds: expected -2<=axis<2.'),
+ exception=(ValueError, errors.InvalidArgumentError)),
])
- def testErrors(self, rt_input, exception, message, axis=1):
+ def testErrors(self, rt_input, exception, message=None, axis=1):
+ rt = ragged_factory_ops.constant(rt_input)
with self.assertRaisesRegexp(exception, message):
- ragged.row_lengths(rt_input, axis)
+ rt.row_lengths(axis)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_row_splits_to_segment_ids_op_test.py b/tensorflow/python/ops/ragged/ragged_row_splits_to_segment_ids_op_test.py
index 7f5f4e9..5384f3a 100644
--- a/tensorflow/python/ops/ragged/ragged_row_splits_to_segment_ids_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_row_splits_to_segment_ids_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for the ragged.row_splits_to_segment_ids() op."""
+"""Tests for the segment_id_ops.row_splits_to_segment_ids() op."""
from __future__ import absolute_import
from __future__ import division
@@ -20,39 +20,36 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.ops.ragged import segment_id_ops
from tensorflow.python.platform import googletest
-class RaggedSplitsToSegmentIdsOpTest(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSplitsToSegmentIdsOpTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testDocStringExample(self):
splits = [0, 3, 3, 5, 6, 9]
expected = [0, 0, 0, 2, 2, 3, 4, 4, 4]
- segment_ids = ragged.row_splits_to_segment_ids(splits)
- with self.test_session():
- self.assertEqual(segment_ids.eval().tolist(), expected)
+ segment_ids = segment_id_ops.row_splits_to_segment_ids(splits)
+ self.assertAllEqual(segment_ids, expected)
- @test_util.run_deprecated_v1
def testEmptySplits(self):
# Note: the splits for an empty ragged tensor contains a single zero.
- segment_ids = ragged.row_splits_to_segment_ids([0])
- with self.test_session():
- self.assertEqual(segment_ids.eval().tolist(), [])
+ segment_ids = segment_id_ops.row_splits_to_segment_ids([0])
+ self.assertAllEqual(segment_ids, [])
- @test_util.run_deprecated_v1
def testErrors(self):
self.assertRaisesRegexp(ValueError, r'Invalid row_splits: \[\]',
- ragged.row_splits_to_segment_ids, [])
+ segment_id_ops.row_splits_to_segment_ids, [])
self.assertRaisesRegexp(
ValueError, r'Tensor conversion requested dtype int64 for '
- 'Tensor with dtype float32', ragged.row_splits_to_segment_ids,
+ 'Tensor with dtype float32', segment_id_ops.row_splits_to_segment_ids,
constant_op.constant([0.5]))
self.assertRaisesRegexp(ValueError, r'Shape \(\) must have rank 1',
- ragged.row_splits_to_segment_ids, 0)
+ segment_id_ops.row_splits_to_segment_ids, 0)
self.assertRaisesRegexp(ValueError, r'Shape \(1, 1\) must have rank 1',
- ragged.row_splits_to_segment_ids, [[0]])
+ segment_id_ops.row_splits_to_segment_ids, [[0]])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_segment_ids_to_row_splits_op_test.py b/tensorflow/python/ops/ragged/ragged_segment_ids_to_row_splits_op_test.py
index 7e52f2d..73ee42a 100644
--- a/tensorflow/python/ops/ragged/ragged_segment_ids_to_row_splits_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_segment_ids_to_row_splits_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for the ragged.segment_ids_to_row_splits() op."""
+"""Tests for the segment_id_ops.segment_ids_to_row_splits() op."""
from __future__ import absolute_import
from __future__ import division
@@ -20,58 +20,52 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.ops.ragged import segment_id_ops
from tensorflow.python.platform import googletest
-class RaggedSplitsToSegmentIdsOpTest(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSplitsToSegmentIdsOpTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testDocStringExample(self):
segment_ids = [0, 0, 0, 2, 2, 3, 4, 4, 4]
expected = [0, 3, 3, 5, 6, 9]
- splits = ragged.segment_ids_to_row_splits(segment_ids)
- with self.test_session():
- self.assertEqual(splits.eval().tolist(), expected)
+ splits = segment_id_ops.segment_ids_to_row_splits(segment_ids)
+ self.assertAllEqual(splits, expected)
- @test_util.run_deprecated_v1
def testEmptySegmentIds(self):
# Note: the splits for an empty ragged tensor contains a single zero.
- segment_ids = ragged.segment_ids_to_row_splits([])
- with self.test_session():
- self.assertEqual(segment_ids.eval().tolist(), [0])
+ segment_ids = segment_id_ops.segment_ids_to_row_splits([])
+ self.assertAllEqual(segment_ids, [0])
def testErrors(self):
self.assertRaisesRegexp(TypeError,
r'segment_ids must be an integer tensor.*',
- ragged.segment_ids_to_row_splits,
+ segment_id_ops.segment_ids_to_row_splits,
constant_op.constant([0.5]))
self.assertRaisesRegexp(ValueError, r'Shape \(\) must have rank 1',
- ragged.segment_ids_to_row_splits, 0)
+ segment_id_ops.segment_ids_to_row_splits, 0)
self.assertRaisesRegexp(ValueError, r'Shape \(1, 1\) must have rank 1',
- ragged.segment_ids_to_row_splits, [[0]])
+ segment_id_ops.segment_ids_to_row_splits, [[0]])
- @test_util.run_deprecated_v1
def testNumSegments(self):
segment_ids = [0, 0, 0, 2, 2, 3, 4, 4, 4]
num_segments = 7
expected = [0, 3, 3, 5, 6, 9, 9, 9]
- splits = ragged.segment_ids_to_row_splits(segment_ids, num_segments)
- with self.test_session():
- self.assertEqual(splits.eval().tolist(), expected)
+ splits = segment_id_ops.segment_ids_to_row_splits(segment_ids, num_segments)
+ self.assertAllEqual(splits, expected)
- @test_util.run_deprecated_v1
def testUnsortedSegmentIds(self):
# Segment ids are not required to be sorted.
segment_ids = [0, 4, 3, 2, 4, 4, 2, 0, 0]
- splits1 = ragged.segment_ids_to_row_splits(segment_ids)
+ splits1 = segment_id_ops.segment_ids_to_row_splits(segment_ids)
expected1 = [0, 3, 3, 5, 6, 9]
- splits2 = ragged.segment_ids_to_row_splits(segment_ids, 7)
+ splits2 = segment_id_ops.segment_ids_to_row_splits(segment_ids, 7)
expected2 = [0, 3, 3, 5, 6, 9, 9, 9]
- with self.test_session():
- self.assertEqual(splits1.eval().tolist(), expected1)
- self.assertEqual(splits2.eval().tolist(), expected2)
+ self.assertAllEqual(splits1, expected1)
+ self.assertAllEqual(splits2, expected2)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_segment_op_test.py b/tensorflow/python/ops/ragged/ragged_segment_op_test.py
index 9e4877a..435ce87 100644
--- a/tensorflow/python/ops/ragged/ragged_segment_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_segment_op_test.py
@@ -25,7 +25,10 @@
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
@@ -45,21 +48,10 @@
return 1.0 * sum(values) / math.sqrt(len(values))
-class RaggedSegmentOpsTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSegmentOpsTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
- def assertNestedListAmostEqual(self, lhs, rhs, places=7, context='value'):
- self.assertEqual(type(lhs), type(rhs))
- if isinstance(lhs, (list, tuple)):
- self.assertEqual(len(lhs), len(rhs), 'Length differs for %s' % context)
- for i in range(len(lhs)):
- self.assertNestedListAmostEqual(lhs[i], rhs[i], places,
- '%s[%s]' % (context, i))
- else:
- self.assertAlmostEqual(
- lhs, rhs, places,
- '%s != %s within %s places at %s' % (lhs, rhs, places, context))
-
def expected_value(self, data, segment_ids, num_segments, combiner):
"""Find the expected value for a call to ragged_segment_<aggregate>.
@@ -71,7 +63,7 @@
Returns:
The expected value, as a nested Python list.
"""
- self.assertEqual(len(data), len(segment_ids))
+ self.assertLen(data, len(segment_ids))
# Build an empty (num_segments x ncols) "grouped" matrix
ncols = max(len(row) for row in data)
@@ -89,147 +81,140 @@
for grouped_row in grouped]
@parameterized.parameters(
- (ragged.segment_sum, sum, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_sum, sum, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_sum, sum, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_sum, sum, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_prod, prod, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_prod, prod, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_prod, prod, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_prod, prod, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_min, min, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_min, min, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_min, min, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_min, min, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_max, max, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_max, max, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_max, max, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_max, max, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_mean, mean, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_mean, mean, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_mean, mean, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_mean, mean, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_sum, sum, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_sum, sum, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_sum, sum, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_sum, sum, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_prod, prod, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_prod, prod, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_prod, prod, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_prod, prod, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_min, min, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_min, min, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_min, min, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_min, min, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_max, max, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_max, max, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_max, max, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_max, max, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_mean, mean, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_mean, mean, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_mean, mean, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_mean, mean, [0, 0, 0, 10, 10, 10]),
)
- @test_util.run_deprecated_v1
def testRaggedSegment_Int(self, segment_op, combiner, segment_ids):
rt_as_list = [[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]]
- rt = ragged.constant(rt_as_list)
+ rt = ragged_factory_ops.constant(rt_as_list)
num_segments = max(segment_ids) + 1
expected = self.expected_value(rt_as_list, segment_ids, num_segments,
combiner)
segmented = segment_op(rt, segment_ids, num_segments)
- self.assertListEqual(self.evaluate(segmented).tolist(), expected)
+ self.assertRaggedEqual(segmented, expected)
@parameterized.parameters(
- (ragged.segment_sum, sum, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_sum, sum, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_sum, sum, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_sum, sum, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_prod, prod, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_prod, prod, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_prod, prod, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_prod, prod, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_min, min, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_min, min, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_min, min, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_min, min, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_max, max, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_max, max, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_max, max, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_max, max, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_mean, mean, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_mean, mean, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_mean, mean, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_mean, mean, [0, 0, 0, 10, 10, 10]),
- (ragged.segment_sqrt_n, sqrt_n, [0, 0, 1, 1, 2, 2]),
- (ragged.segment_sqrt_n, sqrt_n, [0, 0, 0, 1, 1, 1]),
- (ragged.segment_sqrt_n, sqrt_n, [5, 4, 3, 2, 1, 0]),
- (ragged.segment_sqrt_n, sqrt_n, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_sum, sum, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_sum, sum, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_sum, sum, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_sum, sum, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_prod, prod, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_prod, prod, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_prod, prod, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_prod, prod, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_min, min, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_min, min, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_min, min, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_min, min, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_max, max, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_max, max, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_max, max, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_max, max, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_mean, mean, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_mean, mean, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_mean, mean, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_mean, mean, [0, 0, 0, 10, 10, 10]),
+ (ragged_math_ops.segment_sqrt_n, sqrt_n, [0, 0, 1, 1, 2, 2]),
+ (ragged_math_ops.segment_sqrt_n, sqrt_n, [0, 0, 0, 1, 1, 1]),
+ (ragged_math_ops.segment_sqrt_n, sqrt_n, [5, 4, 3, 2, 1, 0]),
+ (ragged_math_ops.segment_sqrt_n, sqrt_n, [0, 0, 0, 10, 10, 10]),
)
- @test_util.run_deprecated_v1
def testRaggedSegment_Float(self, segment_op, combiner, segment_ids):
rt_as_list = [[0., 1., 2., 3.], [4.], [], [5., 6.], [7.], [8., 9.]]
- rt = ragged.constant(rt_as_list)
+ rt = ragged_factory_ops.constant(rt_as_list)
num_segments = max(segment_ids) + 1
expected = self.expected_value(rt_as_list, segment_ids, num_segments,
combiner)
segmented = segment_op(rt, segment_ids, num_segments)
- self.assertNestedListAmostEqual(
- self.evaluate(segmented).tolist(), expected, places=5)
+ self.assertRaggedAlmostEqual(segmented, expected, places=5)
- @test_util.run_deprecated_v1
def testRaggedRankTwo(self):
- rt = ragged.constant([
+ rt = ragged_factory_ops.constant([
[[111, 112, 113, 114], [121],], # row 0
[], # row 1
[[], [321, 322], [331]], # row 2
[[411, 412]] # row 3
]) # pyformat: disable
segment_ids1 = [0, 2, 2, 2]
- segmented1 = ragged.segment_sum(rt, segment_ids1, 3)
+ segmented1 = ragged_math_ops.segment_sum(rt, segment_ids1, 3)
expected1 = [[[111, 112, 113, 114], [121]], # row 0
[], # row 1
[[411, 412], [321, 322], [331]] # row 2
] # pyformat: disable
- self.assertEqual(self.evaluate(segmented1).tolist(), expected1)
+ self.assertRaggedEqual(segmented1, expected1)
segment_ids2 = [1, 2, 1, 1]
- segmented2 = ragged.segment_sum(rt, segment_ids2, 3)
+ segmented2 = ragged_math_ops.segment_sum(rt, segment_ids2, 3)
expected2 = [[],
[[111+411, 112+412, 113, 114], [121+321, 322], [331]],
[]] # pyformat: disable
- self.assertEqual(self.evaluate(segmented2).tolist(), expected2)
+ self.assertRaggedEqual(segmented2, expected2)
- @test_util.run_deprecated_v1
def testRaggedSegmentIds(self):
- rt = ragged.constant([
+ rt = ragged_factory_ops.constant([
[[111, 112, 113, 114], [121],], # row 0
[], # row 1
[[], [321, 322], [331]], # row 2
[[411, 412]] # row 3
]) # pyformat: disable
- segment_ids = ragged.constant([[1, 2], [], [1, 1, 2], [2]])
- segmented = ragged.segment_sum(rt, segment_ids, 3)
+ segment_ids = ragged_factory_ops.constant([[1, 2], [], [1, 1, 2], [2]])
+ segmented = ragged_math_ops.segment_sum(rt, segment_ids, 3)
expected = [[],
[111+321, 112+322, 113, 114],
[121+331+411, 412]] # pyformat: disable
- self.assertEqual(self.evaluate(segmented).tolist(), expected)
+ self.assertRaggedEqual(segmented, expected)
def testShapeMismatchError1(self):
dt = constant_op.constant([1, 2, 3, 4, 5, 6])
- segment_ids = ragged.constant([[1, 2], []])
+ segment_ids = ragged_factory_ops.constant([[1, 2], []])
self.assertRaisesRegexp(
ValueError, 'segment_ids.shape must be a prefix of data.shape, '
- 'but segment_ids is ragged and data is not.', ragged.segment_sum, dt,
- segment_ids, 3)
+ 'but segment_ids is ragged and data is not.',
+ ragged_math_ops.segment_sum, dt, segment_ids, 3)
- @test_util.run_deprecated_v1
def testShapeMismatchError2(self):
- rt = ragged.constant([
+ rt = ragged_factory_ops.constant([
[[111, 112, 113, 114], [121]], # row 0
[], # row 1
[[], [321, 322], [331]], # row 2
[[411, 412]] # row 3
]) # pyformat: disable
- segment_ids = ragged.constant([[1, 2], [1], [1, 1, 2], [2]])
+ segment_ids = ragged_factory_ops.constant([[1, 2], [1], [1, 1, 2], [2]])
# Error is raised at graph-building time if we can detect it then.
self.assertRaisesRegexp(
errors.InvalidArgumentError,
'segment_ids.shape must be a prefix of data.shape.*',
- ragged.segment_sum, rt, segment_ids, 3)
+ ragged_math_ops.segment_sum, rt, segment_ids, 3)
# Otherwise, error is raised when we run the graph.
- segment_ids2 = ragged.from_row_splits(
+ segment_ids2 = ragged_tensor.RaggedTensor.from_row_splits(
array_ops.placeholder_with_default(segment_ids.values, None),
array_ops.placeholder_with_default(segment_ids.row_splits, None))
- segmented2 = ragged.segment_sum(rt, segment_ids2, 3)
- with self.cached_session():
- self.assertRaisesRegexp(
- errors.InvalidArgumentError,
- 'segment_ids.shape must be a prefix of data.shape.*', segmented2.eval)
+ with self.assertRaisesRegexp(
+ errors.InvalidArgumentError,
+ 'segment_ids.shape must be a prefix of data.shape.*'):
+ self.evaluate(ragged_math_ops.segment_sum(rt, segment_ids2, 3))
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_stack_op_test.py b/tensorflow/python/ops/ragged/ragged_stack_op_test.py
index 4343471..f9c8251 100644
--- a/tensorflow/python/ops/ragged/ragged_stack_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_stack_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.stack."""
+"""Tests for ragged_array_ops.stack."""
from __future__ import absolute_import
from __future__ import division
@@ -22,11 +22,15 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedStackOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedStackOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters(
dict(
@@ -265,7 +269,6 @@
axis=0,
expected=[[[b'a00', b'a01'], [], [b'a20', b'a21']]]),
) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedStack(self,
descr,
rt_inputs,
@@ -277,17 +280,16 @@
if ragged_ranks is None:
ragged_ranks = [None] * len(rt_inputs)
rt_inputs = [
- ragged.constant(rt_input, ragged_rank=rrank)
+ ragged_factory_ops.constant(rt_input, ragged_rank=rrank) # pylint: disable=g-long-ternary
if rrank != 0 else constant_op.constant(rt_input)
for (rt_input, rrank) in zip(rt_inputs, ragged_ranks)
]
- stacked = ragged.stack(rt_inputs, axis)
+ stacked = ragged_array_ops.stack(rt_inputs, axis)
if expected_ragged_rank is not None:
self.assertEqual(stacked.ragged_rank, expected_ragged_rank)
if expected_shape is not None:
self.assertEqual(stacked.shape.as_list(), expected_shape)
- with self.test_session():
- self.assertEqual(stacked.eval().tolist(), expected)
+ self.assertRaggedEqual(stacked, expected)
@parameterized.parameters(
dict(
@@ -312,9 +314,9 @@
message='axis=3 out of bounds: expected -3<=axis<3'),
)
def testError(self, rt_inputs, axis, error, message):
- self.assertRaisesRegexp(error, message, ragged.stack, rt_inputs, axis)
+ self.assertRaisesRegexp(error, message, ragged_array_ops.stack, rt_inputs,
+ axis)
- @test_util.run_deprecated_v1
def testSingleTensorInput(self):
"""Tests ragged_stack with a single tensor input.
@@ -322,10 +324,9 @@
also pass in a single value (as with tf.stack), in which case it is
equivalent to expand_dims(axis=0). This test exercises that path.
"""
- rt_inputs = ragged.constant([[1, 2], [3, 4]])
- stacked = ragged.stack(rt_inputs, 0)
- with self.test_session():
- self.assertEqual(stacked.eval().tolist(), [[[1, 2], [3, 4]]])
+ rt_inputs = ragged_factory_ops.constant([[1, 2], [3, 4]])
+ stacked = ragged_array_ops.stack(rt_inputs, 0)
+ self.assertRaggedEqual(stacked, [[[1, 2], [3, 4]]])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_string_ops.py b/tensorflow/python/ops/ragged/ragged_string_ops.py
index cdcdbdf..8021637 100644
--- a/tensorflow/python/ops/ragged/ragged_string_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_string_ops.py
@@ -22,24 +22,27 @@
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_string_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
from tensorflow.python.ops.ragged import ragged_conversion_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.util.tf_export import tf_export
# pylint: disable=redefined-builtin
@tf_export("strings.unicode_encode")
-def unicode_encode(input, output_encoding, errors="replace",
- replacement_char=65533, name=None):
+def unicode_encode(input,
+ output_encoding,
+ errors="replace",
+ replacement_char=65533,
+ name=None):
r"""Encodes each sequence of Unicode code points in `input` into a string.
`result[i1...iN]` is the string formed by concatenating the Unicode
codepoints `input[1...iN, :]`, encoded using `output_encoding`.
Args:
- input: An `N+1` dimensional potentially ragged integer tensor with
- shape `[D1...DN, num_chars]`.
+ input: An `N+1` dimensional potentially ragged integer tensor with shape
+ `[D1...DN, num_chars]`.
output_encoding: Unicode encoding that should be used to encode each
codepoint sequence. Can be `"UTF-8"`, `"UTF-16-BE"`, or `"UTF-32-BE"`.
errors: Specifies the response when an invalid codepoint is encountered
@@ -65,16 +68,16 @@
```
"""
with ops.name_scope(name, "UnicodeEncode", [input]):
- input_tensor = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(input)
+ input_tensor = ragged_tensor.convert_to_tensor_or_ragged_tensor(input)
if input_tensor.shape.ndims is None:
raise ValueError("Rank of input_tensor must be statically known.")
if ragged_tensor.is_ragged(input_tensor):
- if input_tensor.inner_values.shape.ndims > 1:
- # If the inner_values of our ragged tensor is multi-dimensional, we can
+ if input_tensor.flat_values.shape.ndims > 1:
+ # If the flat_values of our ragged tensor is multi-dimensional, we can
# process it separately and our output will have the same nested splits
# as our input.
- return input_tensor.with_inner_values(
- unicode_encode(input_tensor.inner_values, output_encoding, errors,
+ return input_tensor.with_flat_values(
+ unicode_encode(input_tensor.flat_values, output_encoding, errors,
replacement_char))
elif input_tensor.ragged_rank > 1:
# Recursively process the values of the ragged tensor.
@@ -82,7 +85,7 @@
unicode_encode(input_tensor.values, output_encoding, errors,
replacement_char))
else:
- # Our ragged tensor is of the correct shape (rank 1 inner_values tensor
+ # Our ragged tensor is of the correct shape (rank 1 flat_values tensor
# with ragged_rank of 1) so we can process it as normal.
return gen_string_ops.unicode_encode(
input_values=input_tensor.values,
@@ -93,8 +96,9 @@
else:
if input_tensor.shape.ndims == 2:
# The input tensor is of the correct 2-D shape, it's just not ragged.
- return unicode_encode(ragged_conversion_ops.from_tensor(input_tensor),
- output_encoding, errors, replacement_char)
+ return unicode_encode(
+ ragged_conversion_ops.from_tensor(input_tensor), output_encoding,
+ errors, replacement_char)
elif input_tensor.shape.ndims > 2:
# We need to initially flatten the input tensor to 2-D, and then can
# reshape the output of our processed flattened tensor.
@@ -110,10 +114,289 @@
# Our input tensor is rank 1, so we create a ragged tensor with an added
# dimension to create the correct input shape & type, and then remove
# the additional dimension from the output and return the string scalar.
- ragged_input_tensor = ragged_factory_ops.from_row_splits(
+ ragged_input_tensor = ragged_tensor.RaggedTensor.from_row_splits(
input_tensor,
- array_ops.stack([0, array_ops.shape(input_tensor,
- out_type=dtypes.int64)[0]]))
+ array_ops.stack(
+ [0, array_ops.shape(input_tensor, out_type=dtypes.int64)[0]]))
output_tensor = unicode_encode(ragged_input_tensor, output_encoding,
errors, replacement_char)
return array_ops.reshape(output_tensor, [])
+
+
+# pylint: disable=redefined-builtin
+@tf_export("strings.unicode_decode")
+def unicode_decode(input,
+ input_encoding,
+ errors="replace",
+ replacement_char=0xFFFD,
+ replace_control_characters=False,
+ name=None):
+ r"""Decodes each string in `input` into a sequence of Unicode code points.
+
+ `result[i1...iN, j]` is the Unicode codepoint for the `j`th character in
+ `input[i1...iN]`, when decoded using `input_encoding`.
+
+ Args:
+ input: An `N` dimensional potentially ragged `string` tensor with shape
+ `[D1...DN]`. `N` must be statically known.
+ input_encoding: String name for the unicode encoding that should be used to
+ decode each string.
+ errors: Specifies the response when an input string can't be converted
+ using the indicated encoding. One of:
+ * `'strict'`: Raise an exception for any illegal substrings.
+ * `'replace'`: Replace illegal substrings with `replacement_char`.
+ * `'ignore'`: Skip illegal substrings.
+ replacement_char: The replacement codepoint to be used in place of invalid
+ substrings in `input` when `errors='replace'`; and in place of C0 control
+ characters in `input` when `replace_control_characters=True`.
+ replace_control_characters: Whether to replace the C0 control characters
+ `(U+0000 - U+001F)` with the `replacement_char`.
+ name: A name for the operation (optional).
+
+ Returns:
+ A `N+1` dimensional `int32` tensor with shape `[D1...DN, (num_chars)]`.
+ The returned tensor is a `tf.Tensor` if `input` is a scalar, or a
+ `tf.RaggedTensor` otherwise.
+
+ #### Example:
+ ```python
+ >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+ >>> tf.strings.unicode_decode(input, 'UTF-8').tolist()
+ [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]]
+ ```
+ """
+ with ops.name_scope(name, "UnicodeDecode", [input]):
+ return _unicode_decode(input, input_encoding, errors, replacement_char,
+ replace_control_characters, with_offsets=False)
+
+
+@tf_export("strings.unicode_decode_with_offsets")
+def unicode_decode_with_offsets(input,
+ input_encoding,
+ errors="replace",
+ replacement_char=0xFFFD,
+ replace_control_characters=False,
+ name=None):
+ r"""Decodes each string into a sequence of code points with start offsets.
+
+ This op is similar to `tf.strings.decode(...)`, but it also returns the
+ start offset for each character in its respective string. This information
+ can be used to align the characters with the original byte sequence.
+
+ Returns a tuple `(codepoints, start_offsets)` where:
+
+ * `codepoints[i1...iN, j]` is the Unicode codepoint for the `j`th character
+ in `input[i1...iN]`, when decoded using `input_encoding`.
+ * `start_offsets[i1...iN, j]` is the start byte offset for the `j`th
+ character in `input[i1...iN]`, when decoded using `input_encoding`.
+
+ Args:
+ input: An `N` dimensional potentially ragged `string` tensor with shape
+ `[D1...DN]`. `N` must be statically known.
+ input_encoding: String name for the unicode encoding that should be used to
+ decode each string.
+ errors: Specifies the response when an input string can't be converted
+ using the indicated encoding. One of:
+ * `'strict'`: Raise an exception for any illegal substrings.
+ * `'replace'`: Replace illegal substrings with `replacement_char`.
+ * `'ignore'`: Skip illegal substrings.
+ replacement_char: The replacement codepoint to be used in place of invalid
+ substrings in `input` when `errors='replace'`; and in place of C0 control
+ characters in `input` when `replace_control_characters=True`.
+ replace_control_characters: Whether to replace the C0 control characters
+ `(U+0000 - U+001F)` with the `replacement_char`.
+ name: A name for the operation (optional).
+
+ Returns:
+ A tuple of `N+1` dimensional tensors `(codepoints, start_offsets)`.
+
+ * `codepoints` is an `int32` tensor with shape `[D1...DN, (num_chars)]`.
+ * `offsets` is an `int64` tensor with shape `[D1...DN, (num_chars)]`.
+
+ The returned tensors are `tf.Tensor`s if `input` is a scalar, or
+ `tf.RaggedTensor`s otherwise.
+
+ #### Example:
+ ```python
+ >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+ >>> result = tf.strings.unicode_decode_with_offsets(input, 'UTF-8')
+ >>> result[0].tolist() # codepoints
+ [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]]
+ >>> result[1].tolist() # offsets
+ [[0, 1, 3, 5, 6, 7, 8, 9, 10], [0]]
+ ```
+ """
+ with ops.name_scope(name, "UnicodeDecodeWithOffsets", [input]):
+ return _unicode_decode(input, input_encoding, errors, replacement_char,
+ replace_control_characters, with_offsets=True)
+
+
+@tf_export("strings.unicode_split")
+def unicode_split(input,
+ input_encoding,
+ errors="replace",
+ replacement_char=0xFFFD,
+ name=None):
+ r"""Splits each string in `input` into a sequence of Unicode code points.
+
+ `result[i1...iN, j]` is the substring of `input[i1...iN] that encodes its
+ `j`th character, when decoded using `input_encoding`.
+
+ Args:
+ input: An `N` dimensional potentially ragged `string` tensor with shape
+ `[D1...DN]`. `N` must be statically known.
+ input_encoding: String name for the unicode encoding that should be used to
+ decode each string.
+ errors: Specifies the response when an input string can't be converted
+ using the indicated encoding. One of:
+ * `'strict'`: Raise an exception for any illegal substrings.
+ * `'replace'`: Replace illegal substrings with `replacement_char`.
+ * `'ignore'`: Skip illegal substrings.
+ replacement_char: The replacement codepoint to be used in place of invalid
+ substrings in `input` when `errors='replace'`.
+ name: A name for the operation (optional).
+
+ Returns:
+ A `N+1` dimensional `int32` tensor with shape `[D1...DN, (num_chars)]`.
+ The returned tensor is a `tf.Tensor` if `input` is a scalar, or a
+ `tf.RaggedTensor` otherwise.
+
+ #### Example:
+ ```python
+ >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+ >>> tf.strings.unicode_split(input, 'UTF-8').tolist()
+ [['G', '\xc3\xb6', '\xc3\xb6', 'd', 'n', 'i', 'g', 'h', 't'],
+ ['\xf0\x9f\x98\x8a']]
+ ```
+ """
+ with ops.name_scope(name, "UnicodeSplit", [input]):
+ codepoints = _unicode_decode(input, input_encoding, errors,
+ replacement_char, False, with_offsets=False)
+ return unicode_encode(
+ ragged_array_ops.expand_dims(codepoints, -1),
+ output_encoding=input_encoding,
+ errors=errors,
+ replacement_char=replacement_char)
+
+
+@tf_export("strings.unicode_split_with_offsets")
+def unicode_split_with_offsets(input,
+ input_encoding,
+ errors="replace",
+ replacement_char=0xFFFD,
+ name=None):
+ r"""Splits each string into a sequence of code points with start offsets.
+
+ This op is similar to `tf.strings.decode(...)`, but it also returns the
+ start offset for each character in its respective string. This information
+ can be used to align the characters with the original byte sequence.
+
+ Returns a tuple `(chars, start_offsets)` where:
+
+ * `chars[i1...iN, j]` is the substring of `input[i1...iN] that encodes its
+ `j`th character, when decoded using `input_encoding`.
+ * `start_offsets[i1...iN, j]` is the start byte offset for the `j`th
+ character in `input[i1...iN]`, when decoded using `input_encoding`.
+
+ Args:
+ input: An `N` dimensional potentially ragged `string` tensor with shape
+ `[D1...DN]`. `N` must be statically known.
+ input_encoding: String name for the unicode encoding that should be used to
+ decode each string.
+ errors: Specifies the response when an input string can't be converted
+ using the indicated encoding. One of:
+ * `'strict'`: Raise an exception for any illegal substrings.
+ * `'replace'`: Replace illegal substrings with `replacement_char`.
+ * `'ignore'`: Skip illegal substrings.
+ replacement_char: The replacement codepoint to be used in place of invalid
+ substrings in `input` when `errors='replace'`.
+ name: A name for the operation (optional).
+
+ Returns:
+ A tuple of `N+1` dimensional tensors `(codepoints, start_offsets)`.
+
+ * `codepoints` is an `int32` tensor with shape `[D1...DN, (num_chars)]`.
+ * `offsets` is an `int64` tensor with shape `[D1...DN, (num_chars)]`.
+
+ The returned tensors are `tf.Tensor`s if `input` is a scalar, or
+ `tf.RaggedTensor`s otherwise.
+
+ #### Example:
+ ```python
+ >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+ >>> result = tf.strings.unicode_split_with_offsets(input, 'UTF-8')
+ >>> result[0].tolist() # character substrings
+ [['G', '\xc3\xb6', '\xc3\xb6', 'd', 'n', 'i', 'g', 'h', 't'],
+ ['\xf0\x9f\x98\x8a']]
+ >>> result[1].tolist() # offsets
+ [[0, 1, 3, 5, 6, 7, 8, 9, 10], [0]]
+ ```
+ """
+ with ops.name_scope(name, "UnicodeSplitWithOffsets", [input]):
+ codepoints, offsets = _unicode_decode(input, input_encoding, errors,
+ replacement_char, False,
+ with_offsets=True)
+ chars = unicode_encode(
+ ragged_array_ops.expand_dims(codepoints, -1),
+ output_encoding=input_encoding,
+ errors=errors,
+ replacement_char=replacement_char)
+ return chars, offsets
+
+
+def _unicode_decode(input, input_encoding, errors, replacement_char,
+ replace_control_characters, with_offsets):
+ """Decodes each string into a sequence of codepoints."""
+ input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input, name="input")
+ input_ndims = input.shape.ndims
+ if input_ndims is None:
+ raise ValueError("Rank of `input` must be statically known.")
+
+ if input_ndims > 1:
+ # Convert to a ragged tensor with ragged_rank = input_ndims - 1.
+ if not ragged_tensor.is_ragged(input):
+ input = ragged_conversion_ops.from_tensor(
+ input, ragged_rank=input_ndims - 1)
+ elif input.ragged_rank < input_ndims - 1:
+ input = input.with_flat_values(
+ ragged_conversion_ops.from_tensor(
+ input.flat_values,
+ ragged_rank=input_ndims - input.ragged_rank + 1))
+
+ # Reshape the input to a flat vector, and apply the gen_string_ops op.
+ if ragged_tensor.is_ragged(input):
+ flat_input = array_ops.reshape(input.flat_values, [-1])
+ else:
+ flat_input = array_ops.reshape(input, [-1])
+
+ if with_offsets:
+ decode_op = gen_string_ops.unicode_decode_with_offsets
+ else:
+ decode_op = gen_string_ops.unicode_decode
+ flat_result = decode_op(
+ input=flat_input,
+ input_encoding=input_encoding,
+ errors=errors,
+ replacement_char=replacement_char,
+ replace_control_characters=replace_control_characters)
+
+ if input_ndims == 0:
+ codepoints = flat_result.char_values
+ if with_offsets:
+ offsets = flat_result.char_to_byte_starts
+ else:
+ codepoints = ragged_tensor.RaggedTensor.from_row_splits(
+ flat_result.char_values, flat_result.row_splits)
+ if input_ndims > 1:
+ codepoints = input.with_flat_values(codepoints)
+ if with_offsets:
+ offsets = ragged_tensor.RaggedTensor.from_row_splits(
+ flat_result.char_to_byte_starts, flat_result.row_splits)
+ if input_ndims > 1:
+ offsets = input.with_flat_values(offsets)
+
+ if with_offsets:
+ return codepoints, offsets
+ else:
+ return codepoints
+
diff --git a/tensorflow/python/ops/ragged/ragged_tensor.py b/tensorflow/python/ops/ragged/ragged_tensor.py
index ddeabfb..fd334e6 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor.py
@@ -19,9 +19,20 @@
from __future__ import print_function
from tensorflow.python.client import session
+from tensorflow.python.framework import constant_op
+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_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_ragged_conversion_ops
+from tensorflow.python.ops import math_ops
from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.ops.ragged import segment_id_ops
+from tensorflow.python.util.tf_export import tf_export
# pylint: disable=protected-access
_eval_using_default_session = ops._eval_using_default_session
@@ -33,8 +44,9 @@
#===============================================================================
+@tf_export("RaggedTensor")
class RaggedTensor(object):
- """Represents a ragged tensor (go/ragged).
+ """Represents a ragged tensor.
A `RaggedTensor` is a tensor with one or more *ragged dimensions*, which are
dimensions whose slices may have different lengths. For example, the inner
@@ -84,10 +96,10 @@
Example:
```python
- >>> rt = ragged.from_row_splits(values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... row_splits=[0, 4, 4, 7, 8, 8])
- >>> rt.tolist()
- [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
+ >>> print(tf.RaggedTensor.from_row_splits(
+ ... values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... row_splits=[0, 4, 4, 7, 8, 8]))
+ <tf.RaggedTensor [[3, 1, 4, 1], [], [5, 9, 2], [6], []]>
```
### Alternative Row-Partitioning Schemes
@@ -116,13 +128,12 @@
```python
>>> values = [3, 1, 4, 1, 5, 9, 2, 6]
- >>> rt1 = ragged.from_row_splits(values, row_splits=[0, 4, 4, 7, 8, 8])
- >>> rt2 = ragged.from_row_lengths(values, row_lengths=[4, 0, 3, 1, 0])
- >>> rt3 = ragged.from_value_rowids(values,
- ... value_rowids=[0, 0, 0, 0, 2, 2, 2, 3],
- ... nrows=5)
- >>> rt4 = ragged.from_row_starts(values, row_starts=[0, 4, 4, 7, 8])
- >>> rt5 = ragged.from_row_limits(values, row_limits=[4, 4, 7, 8, 8])
+ >>> rt1 = RaggedTensor.from_row_splits(values, row_splits=[0, 4, 4, 7, 8, 8])
+ >>> rt2 = RaggedTensor.from_row_lengths(values, row_lengths=[4, 0, 3, 1, 0])
+ >>> rt3 = RaggedTensor.from_value_rowids(
+ ... values, value_rowids=[0, 0, 0, 0, 2, 2, 2, 3], nrows=5)
+ >>> rt4 = RaggedTensor.from_row_starts(values, row_starts=[0, 4, 4, 7, 8])
+ >>> rt5 = RaggedTensor.from_row_limits(values, row_limits=[4, 4, 7, 8, 8])
```
### Multiple Ragged Dimensions
@@ -132,24 +143,24 @@
adds a single ragged dimension.
```python
- >>> inner_rt = ragged.from_row_splits( # =rt1 from above
+ >>> inner_rt = RaggedTensor.from_row_splits( # =rt1 from above
... values=[3, 1, 4, 1, 5, 9, 2, 6], row_splits=[0, 4, 4, 7, 8, 8])
- >>> outer_rt = ragged.from_row_splits(
+ >>> outer_rt = RaggedTensor.from_row_splits(
... values=inner_rt, row_splits=[0, 3, 3, 5])
- >>> print outer_rt.tolist()
+ >>> print outer_rt.to_list()
[[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]]
>>> print outer_rt.ragged_rank
2
```
- The factory function `ragged.from_nested_row_splits` may be used to
+ The factory function `RaggedTensor.from_nested_row_splits` may be used to
construct a `RaggedTensor` with multiple ragged dimensions directly, by
providing a list of `row_splits` tensors:
```python
- >>> ragged.from_nested_row_splits(
- ... inner_values=[3, 1, 4, 1, 5, 9, 2, 6],
- ... nested_row_splits=([0, 3, 3, 5], [0, 4, 4, 7, 8, 8])).tolist()
+ >>> RaggedTensor.from_nested_row_splits(
+ ... flat_values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... nested_row_splits=([0, 3, 3, 5], [0, 4, 4, 7, 8, 8])).to_list()
[[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]]
```
@@ -159,12 +170,13 @@
by using a multidimensional `Tensor` for `values`.
```python
- >>> rt = ragged.from_row_splits(values=tf.ones([5, 3]), row_splits=[0, 2, 5])
- >>> print rt.tolist()
+ >>> rt = RaggedTensor.from_row_splits(values=tf.ones([5, 3]),
+ .. row_splits=[0, 2, 5])
+ >>> print rt.to_list()
[[[1, 1, 1], [1, 1, 1]],
[[1, 1, 1], [1, 1, 1], [1, 1, 1]]]
- >>> print rt.shape.as_list()
- [2, None, 3]
+ >>> print rt.shape
+ (2, ?, 3)
```
### RaggedTensor Shape Restrictions
@@ -182,31 +194,6 @@
"""
#=============================================================================
- # Implementation notes
- #=============================================================================
- # Currently, the RaggedTensor class uses a single row-partitioning scheme
- # (row_splits).
- #
- # We are considering adding value_rowids+nvals as a secondary
- # row-partitioning scheme. This change would not impact the functional
- # interface of the RaggedTensor class, but it would impact the efficiency
- # of several operations. In particular:
- #
- # * The functions `ragged.value_rowids` and `ragged.nrows` would always
- # return pre-existing tensors; they would not need to add any ops to
- # the graph.
- #
- # * The `RaggedTensor` constructor would construct all row-partitioning
- # tensors (row_splits, value_rowids, and nvals). In eager mode, this
- # would mean that conversion operations would occur whenever a
- # `RaggedTensor` is constructed. But in graph mode, the converted
- # row-partitioning tensors would only be evaluated if they are used.
- #
- # Since this change impacts efficiency but not functionality, we would like
- # to perform additional profiling with real-world use cases before we
- # decide whether to make this change.
-
- #=============================================================================
# Constructor (private)
#=============================================================================
def __init__(self,
@@ -221,13 +208,14 @@
This constructor is private -- please use one of the following ops to
build `RaggedTensor`s:
- * [`ragged.from_row_lengths()`](from_row_lengths.md)
- * [`ragged.from_value_rowids()`](from_value_rowids.md)
- * [`ragged.from_row_splits()`](from_row_splits.md)
- * [`ragged.from_row_starts()`](from_row_starts.md)
- * [`ragged.from_row_limits()`](from_row_limits.md)
- * [`ragged.from_nested_row_splits()`](from_nested_row_splits.md)
- * [`ragged.from_nested_value_rowids()`](from_nested_value_rowids.md)
+ * `tf.RaggedTensor.from_row_lengths`
+ * `tf.RaggedTensor.from_value_rowids`
+ * `tf.RaggedTensor.from_row_splits`
+ * `tf.RaggedTensor.from_row_starts`
+ * `tf.RaggedTensor.from_row_limits`
+ * `tf.RaggedTensor.from_nested_row_splits`
+ * `tf.RaggedTensor.from_nested_row_lengths`
+ * `tf.RaggedTensor.from_nested_value_rowids`
Args:
values: A potentially ragged tensor of any dtype and shape `[nvals, ...]`.
@@ -248,7 +236,7 @@
if not internal:
raise ValueError("RaggedTensor constructor is private; please use one "
"of the factory methods instead (e.g., "
- "ragged.from_row_lengths())")
+ "RaggedTensor.from_row_lengths())")
# Validate the arguments.
if not isinstance(values, (RaggedTensor, ops.Tensor)):
@@ -273,6 +261,364 @@
self._cached_nrows = cached_nrows
#=============================================================================
+ # Factory Methods
+ #=============================================================================
+
+ @classmethod
+ def from_value_rowids(cls, values, value_rowids, nrows=None, name=None):
+ """Creates a `RaggedTensor` with rows partitioned by `value_rowids`.
+
+ The returned `RaggedTensor` corresponds with the python list defined by:
+
+ ```python
+ result = [[values[i] for i in range(len(values)) if value_rowids[i] == row]
+ for row in range(nrows)]
+ ```
+
+ Warning: currently, this needs to cast value_rowids to int64 before
+ converting, since `tf.bincount` only supports `int32`.
+
+ Args:
+ values: A potentially ragged tensor with shape `[nvals, ...]`.
+ value_rowids: A 1-D int64 tensor with shape `[nvals]`, which corresponds
+ one-to-one with `values`, and specifies each value's row index. Must be
+ nonnegative, and must be sorted in ascending order.
+ nrows: An int64 scalar specifying the number of rows. This should be
+ specified if the `RaggedTensor` may containing empty training rows. Must
+ be greater than `value_rowids[-1]` (or zero if `value_rowids` is empty).
+ Defaults to `value_rowids[-1]` (or zero if `value_rowids` is empty).
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor`. `result.rank = values.rank + 1`.
+ `result.ragged_rank = values.ragged_rank + 1`.
+
+ Raises:
+ ValueError: If `nrows` is incompatible with `value_rowids`.
+
+ #### Example:
+ ```python
+ >>> print(tf.RaggedTensor.from_value_rowids(
+ ... values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... value_rowids=[0, 0, 0, 0, 2, 2, 2, 3],
+ ... nrows=5))
+ <tf.RaggedTensor [[3, 1, 4, 1], [], [5, 9, 2], [6], []]>
+ ```
+ """
+ with ops.name_scope(name, "RaggedFromValueRowIds",
+ [values, value_rowids, nrows]):
+ values = convert_to_tensor_or_ragged_tensor(values, name="values")
+ value_rowids = ops.convert_to_tensor(
+ value_rowids, dtypes.int64, name="value_rowids")
+ if nrows is None:
+ const_rowids = tensor_util.constant_value(value_rowids)
+ if const_rowids is None:
+ nrows = array_ops.concat([value_rowids[-1:], [-1]], axis=0)[0] + 1
+ const_nrows = None
+ else:
+ const_nrows = const_rowids[-1] + 1 if const_rowids.size > 0 else 0
+ nrows = ops.convert_to_tensor(const_nrows, dtypes.int64, name="nrows")
+ else:
+ nrows = ops.convert_to_tensor(nrows, dtypes.int64, "nrows")
+ const_nrows = tensor_util.constant_value(nrows)
+ if const_nrows is not None:
+ if const_nrows < 0:
+ raise ValueError("Expected nrows >= 0; got %d" % const_nrows)
+ const_rowids = tensor_util.constant_value(value_rowids)
+ if const_rowids is not None and const_rowids.size > 0:
+ if not const_nrows >= const_rowids[-1] + 1:
+ raise ValueError(
+ "Expected nrows >= value_rowids[-1] + 1; got nrows=%d, "
+ "value_rowids[-1]=%d" % (const_nrows, const_rowids[-1]))
+
+ value_rowids.shape.assert_has_rank(1)
+ nrows.shape.assert_has_rank(0)
+ values.shape[:1].assert_is_compatible_with(value_rowids.shape)
+
+ # Convert value_rowids & nrows to row_splits.
+ # Note: we don't use segment_ids_to_row_splits() here because we want
+ # to save the intermediate value `row_lengths`, so we can cache it.
+ # TODO(b/116708836) Upgrade bincount to accept int64 so we can skip the
+ # cast (Remove the warning in the docstring when we do.)
+ value_rowids_int32 = math_ops.cast(value_rowids, dtypes.int32)
+ nrows_int32 = math_ops.cast(nrows, dtypes.int32)
+ row_lengths = math_ops.bincount(
+ value_rowids_int32,
+ minlength=nrows_int32,
+ maxlength=nrows_int32,
+ dtype=dtypes.int64)
+ row_splits = array_ops.concat([[0], math_ops.cumsum(row_lengths)], axis=0)
+ if const_nrows is not None:
+ row_lengths.set_shape([const_nrows])
+ row_splits.set_shape([const_nrows + 1])
+
+ return cls(
+ values,
+ row_splits,
+ cached_row_lengths=row_lengths,
+ cached_value_rowids=value_rowids,
+ cached_nrows=nrows,
+ internal=True)
+
+ @classmethod
+ def from_row_splits(cls, values, row_splits, name=None):
+ """Creates a `RaggedTensor` with rows partitioned by `row_splits`.
+
+ The returned `RaggedTensor` corresponds with the python list defined by:
+
+ ```python
+ result = [values[row_splits[i]:row_splits[i + 1]]
+ for i in range(len(row_splits) - 1)]
+ ```
+
+ Args:
+ values: A potentially ragged tensor with shape `[nvals, ...]`.
+ row_splits: A 1-D int64 tensor with shape `[nrows+1]`. Must not be empty,
+ and must be sorted in ascending order. `row_splits[0]` must be zero and
+ `row_splits[-1]` must be `nvals`.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor`. `result.rank = values.rank + 1`.
+ `result.ragged_rank = values.ragged_rank + 1`.
+
+ Raises:
+ ValueError: If `row_splits` is an empty list.
+
+ #### Example:
+ ```python
+ >>> print(tf.RaggedTensor.from_row_splits(
+ ... values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... row_splits=[0, 4, 4, 7, 8, 8]))
+ <tf.RaggedTensor [[3, 1, 4, 1], [], [5, 9, 2], [6], []]>
+ ```
+ """
+ if isinstance(row_splits, (list, tuple)) and not row_splits:
+ raise ValueError("row_splits tensor may not be empty.")
+ with ops.name_scope(name, "RaggedFromRowSplits", [values, row_splits]):
+ values = convert_to_tensor_or_ragged_tensor(values, name="values")
+ row_splits = ops.convert_to_tensor(row_splits, dtypes.int64, "row_splits")
+ row_splits.shape.assert_has_rank(1)
+ return cls(values=values, row_splits=row_splits, internal=True)
+
+ @classmethod
+ def from_row_lengths(cls, values, row_lengths, name=None):
+ """Creates a `RaggedTensor` with rows partitioned by `row_lengths`.
+
+ The returned `RaggedTensor` corresponds with the python list defined by:
+
+ ```python
+ result = [[values.pop(0) for i in range(length)]
+ for length in row_lengths]
+ ```
+
+ Args:
+ values: A potentially ragged tensor with shape `[nvals, ...]`.
+ row_lengths: A 1-D int64 tensor with shape `[nrows]`. Must be
+ nonnegative. `sum(row_lengths)` must be `nvals`.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor`. `result.rank = values.rank + 1`.
+ `result.ragged_rank = values.ragged_rank + 1`.
+
+ #### Example:
+ ```python
+ >>> print(tf.RaggedTensor.from_row_lengths(
+ ... values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... row_lengths=[4, 0, 3, 1, 0]))
+ <tf.RaggedTensor [[3, 1, 4, 1], [], [5, 9, 2], [6], []])>
+ ```
+ """
+ with ops.name_scope(name, "RaggedFromRowLengths", [values, row_lengths]):
+ values = convert_to_tensor_or_ragged_tensor(values, name="values")
+ row_lengths = ops.convert_to_tensor(row_lengths, dtypes.int64,
+ "row_lengths")
+ row_lengths.shape.assert_has_rank(1)
+ row_limits = math_ops.cumsum(row_lengths)
+ row_splits = array_ops.concat([[0], row_limits], axis=0)
+ return cls(
+ values=values,
+ row_splits=row_splits,
+ cached_row_lengths=row_lengths,
+ internal=True)
+
+ @classmethod
+ def from_row_starts(cls, values, row_starts, name=None):
+ """Creates a `RaggedTensor` with rows partitioned by `row_starts`.
+
+ Equivalent to: `from_row_splits(values, concat([row_starts, nvals]))`.
+
+ Args:
+ values: A potentially ragged tensor with shape `[nvals, ...]`.
+ row_starts: A 1-D int64 tensor with shape `[nrows]`. Must be nonnegative
+ and sorted in ascending order. If `nrows>0`, then `row_starts[0]` must
+ be zero.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor`. `result.rank = values.rank + 1`.
+ `result.ragged_rank = values.ragged_rank + 1`.
+
+ #### Example:
+ ```python
+ >>> print(tf.RaggedTensor.from_row_starts(
+ ... values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... row_starts=[0, 4, 4, 7, 8]))
+ <tf.RaggedTensor [[3, 1, 4, 1], [], [5, 9, 2], [6], []]>
+ ```
+ """
+ with ops.name_scope(name, "RaggedFromRowStarts", [values, row_starts]):
+ values = convert_to_tensor_or_ragged_tensor(values, name="values")
+ row_starts = ops.convert_to_tensor(row_starts, dtypes.int64, "row_starts")
+ row_starts.shape.assert_has_rank(1)
+ nvals = array_ops.shape(values, out_type=dtypes.int64)[:1]
+ row_splits = array_ops.concat([row_starts, nvals], axis=0)
+ return cls(values=values, row_splits=row_splits, internal=True)
+
+ @classmethod
+ def from_row_limits(cls, values, row_limits, name=None):
+ """Creates a `RaggedTensor` with rows partitioned by `row_limits`.
+
+ Equivalent to: `from_row_splits(values, concat([0, row_limits]))`.
+
+ Args:
+ values: A potentially ragged tensor with shape `[nvals, ...]`.
+ row_limits: A 1-D int64 tensor with shape `[nrows]`. Must be sorted in
+ ascending order. If `nrows>0`, then `row_limits[-1]` must be `nvals`.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor`. `result.rank = values.rank + 1`.
+ `result.ragged_rank = values.ragged_rank + 1`.
+
+ #### Example:
+ ```python
+ >>> print(tf.RaggedTensor.from_row_limits(
+ ... values=[3, 1, 4, 1, 5, 9, 2, 6],
+ ... row_limits=[4, 4, 7, 8, 8]))
+ <tf.RaggedTensor [[3, 1, 4, 1], [], [5, 9, 2], [6], []]>
+ ```
+ """
+ with ops.name_scope(name, "RaggedFromRowLimits", [values, row_limits]):
+ values = convert_to_tensor_or_ragged_tensor(values, name="values")
+ row_limits = ops.convert_to_tensor(row_limits, dtypes.int64, "row_limits")
+ row_limits.shape.assert_has_rank(1)
+ zero = array_ops.zeros([1], dtypes.int64)
+ row_splits = array_ops.concat([zero, row_limits], axis=0)
+ return cls(values=values, row_splits=row_splits, internal=True)
+
+ @classmethod
+ def from_nested_value_rowids(cls,
+ flat_values,
+ nested_value_rowids,
+ nested_nrows=None,
+ name=None):
+ """Creates a `RaggedTensor` from a nested list of `value_rowids` tensors.
+
+ Equivalent to:
+
+ ```python
+ result = flat_values
+ for (rowids, nrows) in reversed(zip(nested_value_rowids, nested_nrows)):
+ result = from_value_rowids(result, rowids, nrows)
+ ```
+
+ Args:
+ flat_values: A potentially ragged tensor.
+ nested_value_rowids: A list of 1-D int64 tensors. The `i`th tensor is
+ used as the `value_rowids` for the `i`th ragged dimension.
+ nested_nrows: A list of int64 scalars. The `i`th scalar is used as the
+ `nrows` for the `i`th ragged dimension.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor` (or `flat_values` if `nested_value_rowids` is empty).
+
+ Raises:
+ ValueError: If `len(nested_values_rowids) != len(nested_nrows)`.
+ """
+ if isinstance(nested_value_rowids, ops.Tensor):
+ raise TypeError("nested_value_rowids must be a list of Tensors")
+ if nested_nrows is None:
+ nested_nrows = [None] * len(nested_value_rowids)
+ else:
+ if isinstance(nested_nrows, ops.Tensor):
+ raise TypeError("nested_nrows must be a list of Tensors")
+ if len(nested_nrows) != len(nested_value_rowids):
+ raise ValueError("nested_nrows must have the same length as "
+ "nested_value_rowids")
+
+ with ops.name_scope(
+ name, "RaggedFromNestedValueRowIds",
+ [flat_values] + list(nested_value_rowids) + list(nested_nrows)):
+ result = flat_values
+ for value_rowids, nrows in reversed(
+ list(zip(nested_value_rowids, nested_nrows))):
+ result = cls.from_value_rowids(result, value_rowids, nrows)
+ return result
+
+ @classmethod
+ def from_nested_row_splits(cls, flat_values, nested_row_splits, name=None):
+ """Creates a `RaggedTensor` from a nested list of `row_splits` tensors.
+
+ Equivalent to:
+
+ ```python
+ result = flat_values
+ for row_splits in reversed(nested_row_splits):
+ result = from_row_splits(result, row_splits)
+ ```
+
+ Args:
+ flat_values: A potentially ragged tensor.
+ nested_row_splits: A list of 1-D int64 tensors. The `i`th tensor is used
+ as the `row_splits` for the `i`th ragged dimension.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor` (or `flat_values` if `nested_row_splits` is empty).
+ """
+ if isinstance(nested_row_splits, ops.Tensor):
+ raise TypeError("nested_row_splits must be a list of Tensors")
+ with ops.name_scope(name, "RaggedFromNestedRowSplits",
+ [flat_values] + list(nested_row_splits)):
+ result = flat_values
+ for splits in reversed(nested_row_splits):
+ result = cls.from_row_splits(result, splits)
+ return result
+
+ @classmethod
+ def from_nested_row_lengths(cls, flat_values, nested_row_lengths, name=None):
+ """Creates a `RaggedTensor` from a nested list of `row_lengths` tensors.
+
+ Equivalent to:
+
+ ```python
+ result = flat_values
+ for row_lengths in reversed(nested_row_lengths):
+ result = from_row_lengths(result, row_lengths)
+ ```
+
+ Args:
+ flat_values: A potentially ragged tensor.
+ nested_row_lengths: A list of 1-D int64 tensors. The `i`th tensor is used
+ as the `row_lengths` for the `i`th ragged dimension.
+ name: A name prefix for the RaggedTensor (optional).
+
+ Returns:
+ A `RaggedTensor` (or `flat_values` if `nested_row_lengths` is empty).
+ """
+ if isinstance(nested_row_lengths, ops.Tensor):
+ raise TypeError("nested_row_lengths must be a list of Tensors")
+ with ops.name_scope(name, "RaggedFromNestedRowlengths",
+ [flat_values] + list(nested_row_lengths)):
+ result = flat_values
+ for lengths in reversed(nested_row_lengths):
+ result = cls.from_row_lengths(result, lengths)
+ return result
+
+ #=============================================================================
# Accessors
#=============================================================================
@@ -334,8 +680,8 @@
#### Example:
```python
>>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
- >>> rt.values.eval()
- [3, 1, 4, 1, 5, 9, 2, 6]
+ >>> print rt.values
+ tf.Tensor([3, 1, 4, 1, 5, 9, 2, 6])
```
"""
return self._values
@@ -357,26 +703,24 @@
#### Example:
```python
>>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
- >>> rt.values.eval()
- [3, 1, 4, 1, 5, 9, 2, 6]
- >>> rt.row_splits.eval() # indices of row splits in ragged.values
- [0, 4, 4, 7, 8, 8]
+ >>> print rt.row_splits # indices of row splits in rt.values
+ tf.Tensor([0, 4, 4, 7, 8, 8])
```
"""
return self._row_splits
@property
- def inner_values(self):
+ def flat_values(self):
"""The innermost `values` tensor for this ragged tensor.
- Concretely, if `rt.values` is a `Tensor`, then `rt.inner_values` is
- `rt.values`; otherwise, `rt.inner_values` is `rt.values.inner_values`.
+ Concretely, if `rt.values` is a `Tensor`, then `rt.flat_values` is
+ `rt.values`; otherwise, `rt.flat_values` is `rt.values.flat_values`.
- Conceptually, `inner_values` is the tensor formed by flattening the
+ Conceptually, `flat_values` is the tensor formed by flattening the
outermost dimension and all of the ragged dimensions into a single
dimension.
- `rt.inner_values.shape = [nvals] + rt.shape[rt.ragged_rank + 1:]`
+ `rt.flat_values.shape = [nvals] + rt.shape[rt.ragged_rank + 1:]`
(where `nvals` is the number of items in the flattened dimensions).
Returns:
@@ -386,8 +730,8 @@
```python
>>> rt = ragged.constant([[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]])
- >>> ragged.inner_values(rt).eval()
- [3, 1, 4, 1, 5, 9, 2, 6]
+ >>> print rt.flat_values()
+ tf.Tensor([3, 1, 4, 1, 5, 9, 2, 6])
```
"""
rt_values = self.values
@@ -413,8 +757,8 @@
```python
>>> rt = ragged.constant([[[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]]])
- >>> for i, splits in enumerate(ragged.nested_row_splits(rt)):
- ... print('Splits for dimension %d: %s' % (i+1, splits.eval()))
+ >>> for i, splits in enumerate(rt.nested_row_splits()):
+ ... print('Splits for dimension %d: %s' % (i+1, splits))
Splits for dimension 1: [0, 1]
Splits for dimension 2: [0, 3, 3, 5]
Splits for dimension 3: [0, 4, 4, 7, 8, 8]
@@ -428,38 +772,220 @@
rt_values = rt_values.values
return tuple(rt_nested_splits)
- @property
- def cached_value_rowids(self):
- """The row lengths for this `RaggedTensor`, or `None`.
+ def value_rowids(self, name=None):
+ """Returns the row indices for the `values` in this ragged tensor.
+
+ `rt.value_rowids()` corresponds one-to-one with the outermost dimension of
+ `rt.values`, and specifies the row containing each value. In particular,
+ the row `rt[row]` consists of the values `rt.values[j]` where
+ `rt.value_rowids()[j] == row`.
+
+ Args:
+ name: A name prefix for the returned tensor (optional).
Returns:
- The `value_rowids` tensor that was used to construct this `RaggedTensor`
- if it was constructed using
- [`ragged.from_value_rowids`](from_value_rowids.md); or `None` otherwise.
- """
- return self._cached_value_rowids
+ A 1-D `int64` `Tensor` with shape `self.values.shape[:1]`.
+ The returned tensor is nonnegative, and is sorted in ascending order.
- @property
- def cached_nrows(self):
- """The row lengths for this `RaggedTensor`, or `None`.
+ #### Example:
+ ```python
+ >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
+ >>> rt.values
+ tf.Tensor([3, 1, 4, 1, 5, 9, 2, 6])
+ >>> rt.value_rowids()
+ tf.Tensor([0, 0, 0, 0, 2, 2, 2, 3]) # corresponds 1:1 with rt.values
+ ```
+ """
+ if self._cached_value_rowids is not None:
+ return self._cached_value_rowids
+
+ with ops.name_scope(name, "RaggedValueRowIds", [self]):
+ return segment_id_ops.row_splits_to_segment_ids(self.row_splits)
+
+ def nrows(self, out_type=dtypes.int64, name=None):
+ """Returns the number of rows in this ragged tensor.
+
+ I.e., the size of the outermost dimension of the tensor.
+
+ Args:
+ out_type: `dtype` for the returned tensor.
+ name: A name prefix for the returned tensor (optional).
Returns:
- The `nrows` tensor that was used to construct this `RaggedTensor`
- if it was constructed using
- [`ragged.from_value_rowids`](from_value_rowids.md); or `None` otherwise.
- """
- return self._cached_nrows
+ A scalar `Tensor` with dtype `out_type`.
- @property
- def cached_row_lengths(self):
- """The row lengths for this `RaggedTensor`, or `None`.
+ #### Example:
+ ```python
+ >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
+ >>> rt.nrows() # rt has 5 rows.
+ 5
+ ```
+ """
+ if self._cached_nrows is not None:
+ return self._cached_nrows
+
+ with ops.name_scope(name, "RaggedNRows", [self]):
+ return array_ops.shape(self.row_splits, out_type=out_type)[0] - 1
+
+ def row_starts(self, name=None):
+ """Returns the start indices for rows in this ragged tensor.
+
+ These indices specify where the values for each row begin in
+ `self.values`. `rt.row_starts()` is equal to `rt.row_splits[:-1]`.
+
+ Args:
+ name: A name prefix for the returned tensor (optional).
Returns:
- The `row_lengths` tensor that was used to construct this `RaggedTensor`
- if it was constructed using
- [`ragged.from_row_lengths`](from_row_lengths.md); or `None` otherwise.
+ A 1-D Tensor of int64 with shape `[nrows]`.
+ The returned tensor is nonnegative, and is sorted in ascending order.
+
+ #### Example:
+ ```python
+ >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
+ >>> rt.values
+ tf.Tensor([3, 1, 4, 1, 5, 9, 2, 6])
+ >>> rt.row_starts() # indices of row starts in rt.values
+ tf.Tensor([0, 4, 4, 7, 8])
+ ```
"""
- return self._cached_row_lengths
+ with ops.name_scope(name, "RaggedRowStarts", [self]):
+ return self.row_splits[:-1]
+
+ def row_limits(self, name=None):
+ """Returns the limit indices for rows in this ragged tensor.
+
+ These indices specify where the values for each row end in
+ `self.values`. `rt.row_limits(self)` is equal to `rt.row_splits[:-1]`.
+
+ Args:
+ name: A name prefix for the returned tensor (optional).
+
+ Returns:
+ A 1-D Tensor of int64 with shape `[nrows]`.
+ The returned tensor is nonnegative, and is sorted in ascending order.
+
+ #### Example:
+ ```python
+ >>> rt = ragged.constant([[3, 1, 4, 1], [], [5, 9, 2], [6], []])
+ >>> rt.values
+ tf.Tensor([3, 1, 4, 1, 5, 9, 2, 6])
+ >>> rt.row_limits() # indices of row limits in rt.values
+ tf.Tensor([4, 4, 7, 8, 8])
+ ```
+ """
+ with ops.name_scope(name, "RaggedRowLimits", [self]):
+ return self.row_splits[1:]
+
+ def row_lengths(self, axis=1, name=None):
+ """Returns the lengths of the rows in this ragged tensor.
+
+ `rt.row_lengths()[i]` indicates the number of values in the
+ `i`th row of `rt`.
+
+ Args:
+ axis: An integer constant indicating the axis whose row lengths should be
+ returned.
+ name: A name prefix for the returned tensor (optional).
+
+ Returns:
+ A potentially ragged Tensor of int64 with shape `self.shape[:axis]`.
+
+ Raises:
+ ValueError: If `axis` is out of bounds.
+
+ #### Example:
+ ```python
+ >>> rt = ragged.constant([[[3, 1, 4], [1]], [], [[5, 9], [2]], [[6]], []])
+ >>> rt.row_lengths(rt) # lengths of rows in rt
+ tf.Tensor([2, 0, 2, 1, 0])
+ >>> rt.row_lengths(axis=2) # lengths of axis=2 rows.
+ <tf.RaggedTensor [[3, 1], [], [2, 1], [1], []]>
+ ```
+ """
+ if self._cached_row_lengths is not None:
+ return self._cached_row_lengths
+
+ with ops.name_scope(name, "RaggedRowLengths", [self]):
+ axis = ragged_util.get_positive_axis(axis, self.shape.ndims)
+ if axis == 0:
+ return self.nrows()
+ elif axis == 1:
+ splits = self.row_splits
+ return splits[1:] - splits[:-1]
+ elif isinstance(self.values, RaggedTensor):
+ return self.with_values(self.values.row_lengths(axis - 1))
+ else:
+ shape = array_ops.shape(self.values, out_type=dtypes.int64)
+ return self.with_values(
+ array_ops.ones(shape[:axis - 1], dtypes.int64) * shape[axis - 1])
+
+ def nested_row_lengths(self, name=None):
+ """Returns a tuple containing the row_lengths for all ragged dimensions.
+
+ `rtnested_row_lengths()` is a tuple containing the `row_lengths` tensors for
+ all ragged dimensions in `rt`, ordered from outermost to innermost.
+
+ Args:
+ name: A name prefix for the returned tensors (optional).
+
+ Returns:
+ A `tuple` of 1-D `int64` `Tensors`. The length of the tuple is equal to
+ `self.ragged_rank`.
+ """
+ with ops.name_scope(name, "RaggedNestedRowLengths", [self]):
+ rt_nested_row_lengths = []
+ rt = self
+ while isinstance(rt, RaggedTensor):
+ rt_nested_row_lengths.append(rt.row_lengths())
+ rt = rt.values
+ return tuple(rt_nested_row_lengths)
+
+ def bounding_shape(self, axis=None, name=None):
+ """Returns the tight bounding box shape for this `RaggedTensor`.
+
+ Args:
+ axis: An integer scalar or vector indicating which axes to return the
+ bounding box for. If not specified, then the full bounding box is
+ returned.
+ name: A name prefix for the returned tensor (optional).
+
+ Returns:
+ An int64 `Tensor`. If `axis` is not specified, then `output`
+ is a vector with `output.shape=[self.shape.ndims]`. If `axis` is a
+ scalar, then the `output` is a scalar. If `axis` is a vector, then
+ `output` is a vector, where `output[i]` is the bounding size for
+ dimension `axis[i]`.
+
+ #### Example:
+ ```python
+ >>> rt = ragged.constant([[1, 2, 3, 4], [5], [], [6, 7, 8, 9], [10]])
+ >>> rt.bounding_shape()
+ [5, 4]
+ ```
+ """
+ with ops.name_scope(name, "RaggedBoundingBox", [self, axis]):
+ nested_splits = self.nested_row_splits
+ rt_flat_values = self.flat_values
+
+ # Optimized special cases for when axis=0 or axis=1:
+ if isinstance(axis, int):
+ if axis == 0:
+ return array_ops.shape(nested_splits[0], out_type=dtypes.int64)[0] - 1
+ elif axis == 1:
+ return math_ops.maximum(math_ops.reduce_max(self.row_lengths()), 0)
+
+ splits_shape = array_ops.shape(self.row_splits, out_type=dtypes.int64)
+ flat_values_shape = array_ops.shape(rt_flat_values, out_type=dtypes.int64)
+
+ ragged_dimensions = array_ops.stack([splits_shape[0] - 1] + [
+ math_ops.maximum(math_ops.reduce_max(splits[1:] - splits[:-1]), 0)
+ for splits in nested_splits
+ ])
+ inner_dimensions = flat_values_shape[1:]
+
+ bbox = array_ops.concat([ragged_dimensions, inner_dimensions], axis=0)
+ return bbox if axis is None else array_ops.gather(bbox, axis)
#=============================================================================
# Transformation
@@ -481,7 +1007,7 @@
`result.ragged_rank = 1 + new_values.ragged_rank`
"""
new_values.shape.with_rank_at_least(1)
- self.values.shape[0].assert_is_compatible_with(new_values.shape[0])
+ self.values.shape[:1].assert_is_compatible_with(new_values.shape[:1])
return RaggedTensor(
new_values,
self._row_splits,
@@ -490,16 +1016,16 @@
self._cached_nrows,
internal=True)
- def with_inner_values(self, new_values):
- """Returns a copy of `self` with `inner_values` replaced by `new_value`.
+ def with_flat_values(self, new_values):
+ """Returns a copy of `self` with `flat_values` replaced by `new_value`.
Preserves cached row-partitioning tensors such as `self.cached_nrows` and
`self.cached_value_rowids` if they have values.
Args:
new_values: Potentially ragged tensor that should replace
- `self.inner_values`. Must have `rank > 0`, and must have the same
- number of rows as `self.inner_values`.
+ `self.flat_values`. Must have `rank > 0`, and must have the same
+ number of rows as `self.flat_values`.
Returns:
A `RaggedTensor`.
@@ -509,46 +1035,369 @@
if isinstance(self._values, ops.Tensor):
return self.with_values(new_values)
else:
- return self.with_values(self.values.with_inner_values(new_values))
+ return self.with_values(self.values.with_flat_values(new_values))
+
+ #=============================================================================
+ # Tensor Type Conversions
+ #=============================================================================
+
+ @classmethod
+ def from_tensor(cls,
+ tensor,
+ lengths=None,
+ padding=None,
+ ragged_rank=1,
+ name=None):
+ """Converts a `tf.Tensor` into a `RaggedTensor`.
+
+ The set of absent/default values may be specified using a vector of lengths
+ or a padding value (but not both). If `lengths` is specified, then the
+ output tensor will satisfy `output[row] = tensor[row][:lengths[row]]`.
+ If `padding` is specified, then any row *suffix* consisting entirely of
+ `padding` will be excluded from the returned `RaggedTensor`. If neither
+ `lengths` nor `padding` is specified, then the returned `RaggedTensor` will
+ have no absent/default values.
+
+ Examples:
+
+ ```python
+ >>> dt = tf.constant([[5, 7, 0], [0, 3, 0], [6, 0, 0]])
+ >>> tf.RaggedTensor.from_tensor(dt)
+ <tf.RaggedTensor [[5, 7, 0], [0, 3, 0], [6, 0, 0]]>
+ >>> tf.RaggedTensor.from_tensor(dt, lengths=[2, 0, 3])
+ <tf.RaggedTensor [[5, 7], [], [6, 0, 0]]>
+ >>> tf.RaggedTensor.from_tensor(dt, padding=0)
+ <tf.RaggedTensor [[5, 7], [0, 3], [6]]>
+ ```
+
+ Args:
+ tensor: The `Tensor` to convert. Must have rank `ragged_rank + 1` or
+ higher.
+ lengths: An optional set of row lengths, specified using a 1-D integer
+ `Tensor` whose length is equal to `tensor.shape[0]` (the number of rows
+ in `tensor`). If specified, then `output[row]` will contain
+ `tensor[row][:lengths[row]]`. Negative lengths are treated as zero.
+ padding: An optional padding value. If specified, then any row suffix
+ consisting entirely of `padding` will be excluded from the returned
+ RaggedTensor. `padding` is a `Tensor` with the same dtype as `tensor`
+ and with `shape=tensor.shape[ragged_rank + 1:]`.
+ ragged_rank: Integer specifying the ragged rank for the returned
+ `RaggedTensor`. Must be greater than zero.
+ name: A name prefix for the returned tensors (optional).
+
+ Returns:
+ A `RaggedTensor` with the specified `ragged_rank`. The shape of the
+ returned ragged tensor is compatible with the shape of `tensor`.
+ Raises:
+ ValueError: If both `lengths` and `padding` are specified.
+ """
+ if lengths is not None and padding is not None:
+ raise ValueError("Specify lengths or padding, but not both")
+ if not isinstance(ragged_rank, int):
+ raise TypeError("ragged_rank expected int, got %r" % ragged_rank)
+ if ragged_rank <= 0:
+ raise ValueError(
+ "ragged_rank must be greater than 0; got %s" % ragged_rank)
+
+ with ops.name_scope(name, "RaggedFromTensor", [tensor, lengths, padding]):
+ tensor = ops.convert_to_tensor(tensor, name="tensor")
+ tensor.shape.with_rank_at_least(ragged_rank + 1)
+ input_shape = array_ops.shape(tensor, out_type=dtypes.int64)
+ ncols = input_shape[1]
+
+ # Handle ragged_rank>1 via recursion:
+ # If the output should have multiple ragged dimensions, then first
+ # flatten the tensor to eliminate all but the last ragged dimension,
+ # and recursively convert that flattened tensor. Then add on the splits
+ # for the dimensions that we flattened out.
+ if ragged_rank > 1:
+ # Flatten `tensor` to eliminate all but the last ragged dimension.
+ new_shape = array_ops.concat([
+ constant_op.constant([-1], dtypes.int64), input_shape[ragged_rank:]
+ ],
+ axis=0)
+ flattened = array_ops.reshape(tensor, new_shape)
+ # Recursively convert the flattened tensor.
+ values = cls.from_tensor(flattened, lengths, padding)
+ # The total number of elements in each dimension. E.g., if
+ # input_shape=[3, 4, 5, 6], then dim[2] has 3*4*5 elements in total.
+ dim_size = math_ops.cumprod(input_shape)
+ # Construct splits tensors for the dimensions that were flattened.
+ new_splits = [
+ math_ops.range(0, dim_size[dim - 1] + 1) * input_shape[dim]
+ for dim in range(1, ragged_rank)
+ ]
+ return cls.from_nested_row_splits(values, new_splits)
+
+ # If padding was specified, then use it to find row lengths.
+ if padding is not None:
+ padding = ops.convert_to_tensor(
+ padding, name="padding", dtype=tensor.dtype)
+ padding.shape.assert_is_compatible_with(tensor.shape[2:])
+
+ # Find places where the padding is equal to the tensor. (This will
+ # broadcast `padding` across the outermost 2 dimensions of `tensor`,
+ # so `has_default_value.shape = tensor.shape`.)
+ has_default_value = math_ops.equal(padding, tensor)
+
+ # If the padding isn't a scalar, then require that all values in the
+ # padding match each item in the tensor. After this block of code,
+ # `has_default.shape = tensor.shape[:2]`. (Unfortunately, we can't just
+ # use reduce_all for both cases, becaue when you pass an empty `axis`
+ # list to reduce_all, it reduces all axes; but we want it to reduce no
+ # axes -- i.e., to be a no-op.)
+ tensor_rank = array_ops.rank(tensor)
+ reduce_axis = math_ops.range(2, tensor_rank)
+ has_default = control_flow_ops.cond(
+ tensor_rank > 2,
+ lambda: math_ops.reduce_all(has_default_value, axis=reduce_axis),
+ lambda: has_default_value)
+ has_default.set_shape(tensor_shape.TensorShape([None, None]))
+ has_default.set_shape(tensor.shape[:2])
+
+ # Use has_default it to find the length of each row: for each
+ # non-default item in a row, calculate the length that the row needs to
+ # have to include that item; and then take the max of those values
+ # (across each row).
+ has_nondefault = math_ops.logical_not(has_default)
+ has_nondefault = math_ops.cast(has_nondefault, dtypes.int64)
+ length_for_nondefault_value = (
+ has_nondefault * array_ops.expand_dims(
+ math_ops.range(1, ncols + 1), 0))
+ lengths = math_ops.reduce_max(length_for_nondefault_value, axis=1)
+
+ # If we have lengths (either directly supplied, or computed from
+ # paddings), then use those to construct splits; and then use masking
+ # to get the corresponding values.
+ if lengths is not None:
+ lengths = ragged_util.convert_to_int_tensor(lengths, "lengths",
+ dtypes.int64)
+ lengths.shape.assert_has_rank(1)
+ lengths = math_ops.minimum(lengths, ncols)
+ lengths = math_ops.maximum(lengths, 0)
+ limits = math_ops.cumsum(lengths)
+ splits = array_ops.concat([array_ops.zeros([1], dtypes.int64), limits],
+ axis=0)
+ mask = array_ops.sequence_mask(lengths, maxlen=ncols)
+ values = array_ops.boolean_mask(tensor, mask)
+ return cls.from_row_splits(values, splits)
+
+ # If neither padding nor lengths were specified, then create a splits
+ # vector that contains no default values, and reshape the input tensor
+ # to form the values for the RaggedTensor.
+ nrows = input_shape[0]
+ nvals = nrows * ncols
+ splits = math_ops.range(nrows + 1) * ncols
+ values_shape = array_ops.concat([[nvals], input_shape[2:]], axis=0)
+ values = array_ops.reshape(tensor, values_shape)
+ return cls.from_row_splits(values, splits)
+
+ def to_tensor(self, default_value=None, name=None):
+ """Converts this `RaggedTensor` into a `tf.Tensor`.
+
+ Example:
+
+ ```python
+ >>> rt = ragged.constant([[9, 8, 7], [], [6, 5], [4]])
+ >>> print rt.to_tensor()
+ [[9 8 7]
+ [0 0 0]
+ [6 5 0]
+ [4 0 0]]
+ ```
+
+ Args:
+ default_value: Value to set for indices not specified in `self`. Defaults
+ to zero. `default_value` must be broadcastable to
+ `self.shape[self.ragged_rank + 1:]`.
+ name: A name prefix for the returned tensors (optional).
+
+ Returns:
+ A `Tensor` with shape `ragged.bounding_shape(self)` and the
+ values specified by the non-empty values in `self`. Empty values are
+ assigned `default_value`.
+ """
+ with ops.name_scope(name, "RaggedToTensor", [self, default_value]):
+ if default_value is not None:
+ default_value = ops.convert_to_tensor(
+ default_value, name="default_value", dtype=self.dtype)
+
+ # If ragged_rank > 1, then recursively convert the ragged values into a
+ # `Tensor` before we proceed.
+ values = self.values
+ if is_ragged(values):
+ values = values.to_tensor(default_value)
+
+ # Tile the default value, if necessary.
+ if default_value is not None:
+ if values.shape.ndims is not None:
+ default_value.shape.with_rank_at_most(values.shape.ndims - 1)
+ if (values.shape.ndims is None or default_value.shape.ndims is None or
+ values.shape.ndims != default_value.shape.ndims + 1):
+ value_shape = array_ops.shape(values)[1:]
+ default_value = array_ops.broadcast_to(default_value, value_shape)
+ default_value.shape.assert_is_compatible_with(values.shape[1:])
+
+ # Get the expected dense shape ([nrows, ncols] + value_shape).
+ rt_row_lengths = [self.row_splits[1:] - self.row_splits[:-1]]
+ nrows = array_ops.shape(self.row_splits, out_type=dtypes.int64)[0] - 1
+ ncols = math_ops.maximum(math_ops.reduce_max(rt_row_lengths), 0)
+ values_shape = array_ops.shape(values, out_type=dtypes.int64)
+ value_shape = values_shape[1:]
+ nvals = values_shape[0]
+
+ # Build a default value if none was supplied.
+ if default_value is None:
+ default_value = array_ops.zeros(value_shape, dtype=values.dtype)
+ default_value.shape.assert_is_compatible_with(values.shape[1:])
+ default_value.set_shape(values.shape[1:])
+
+ # Get the row start indices, and expand to shape=[nrows, 1].
+ starts = array_ops.expand_dims(self.row_splits[:-1], 1)
+
+ # Get the row limit indices, and expand to shape=[nrows, 1].
+ limits = array_ops.expand_dims(self.row_splits[1:], 1)
+
+ # Get the column indices, and expand to shape=[1, ncols].
+ columns = array_ops.expand_dims(math_ops.range(0, ncols), 0)
+
+ # Build a list containing the values plus the default value. We will use
+ # tf.gather to collect values from this list for the `Tensor` (using
+ # nvals as the index for the default value).
+ values_and_default = array_ops.concat(
+ [values, array_ops.stack([default_value])], axis=0)
+
+ # Construct a matrix "indices" pointing into values_and_default. I.e.,
+ # output[r, c] = values_and_default[indices[r, c].
+ nondefault_index = starts + columns
+ has_value = nondefault_index < limits
+ default_index = array_ops.fill(array_ops.stack([nrows, ncols]), nvals)
+ indices = array_ops.where(has_value, nondefault_index, default_index)
+
+ # Gather the results into a `Tensor`.
+ return array_ops.gather(values_and_default, indices)
+
+ @classmethod
+ def from_sparse(cls, st_input, name=None):
+ """Converts a 2D `tf.SparseTensor` to a `RaggedTensor`.
+
+ Each row of the `output` `RaggedTensor` will contain the explicit values
+ from the same row in `st_input`. `st_input` must be ragged-right. If not
+ it is not ragged-right, then an error will be generated.
+
+ Example:
+
+ ```python
+ >>> st = SparseTensor(indices=[[0, 1], [0, 2], [0, 3], [1, 0], [3, 0]],
+ ... values=[1, 2, 3, 4, 5],
+ ... dense_shape=[4, 3])
+ >>> rt.RaggedTensor.from_sparse(st).eval().tolist()
+ [[1, 2, 3], [4], [], [5]]
+ ```
+
+ Currently, only two-dimensional `SparseTensors` are supported.
+
+ Args:
+ st_input: The sparse tensor to convert. Must have rank 2.
+ name: A name prefix for the returned tensors (optional).
+
+ Returns:
+ A `RaggedTensor` with the same values as `st_input`.
+ `output.ragged_rank = rank(st_input) - 1`.
+ `output.shape = [st_input.dense_shape[0], None]`.
+ Raises:
+ ValueError: If the number of dimensions in `st_input` is not known
+ statically, or is not two.
+ """
+ if not sparse_tensor.is_sparse(st_input):
+ raise TypeError("Expected SparseTensor, got %s" % type(st_input).__name__)
+ with ops.name_scope(name, "RaggedFromSparse", [st_input]):
+ st_input = sparse_tensor.convert_to_tensor_or_sparse_tensor(
+ st_input, name="st_input")
+
+ if st_input.dense_shape.shape.ndims is None:
+ static_rank_from_dense_shape = None
+ else:
+ static_rank_from_dense_shape = st_input.dense_shape.shape.dims[0].value
+
+ if st_input.indices.shape.ndims is None:
+ static_rank_from_indices = None
+ else:
+ static_rank_from_indices = st_input.indices.shape.dims[1].value
+
+ if static_rank_from_dense_shape != 2 and static_rank_from_indices != 2:
+ raise ValueError("rank(st_input) must be 2")
+
+ with ops.control_dependencies(
+ _assert_sparse_indices_are_ragged_right(st_input.indices)):
+ # Treat sparse row indices as segment ids to generate a splits tensor
+ # thta we can pair with the sparse tensor values. (Ignore sparse column
+ # indices.)
+ segment_ids = st_input.indices[:, 0]
+ num_segments = st_input.dense_shape[0]
+ return cls.from_value_rowids(st_input.values, segment_ids, num_segments)
+
+ def to_sparse(self, name=None):
+ """Converts this `RaggedTensor` into a `tf.SparseTensor`.
+
+ Example:
+
+ ```python
+ >>> rt = ragged.constant([[1, 2, 3], [4], [], [5, 6]])
+ >>> rt.to_sparse().eval()
+ SparseTensorValue(indices=[[0, 0], [0, 1], [0, 2], [1, 0], [3, 0], [3, 1]],
+ values=[1, 2, 3, 4, 5, 6],
+ dense_shape=[4, 3])
+ ```
+
+ Args:
+ name: A name prefix for the returned tensors (optional).
+
+ Returns:
+ A SparseTensor with the same values as `self`.
+ """
+ with ops.name_scope(name, "RaggedToSparse", [self]):
+ result = gen_ragged_conversion_ops.ragged_tensor_to_sparse(
+ self.nested_row_splits, self.flat_values, name=name)
+ return sparse_tensor.SparseTensor(result.sparse_indices,
+ result.sparse_values,
+ result.sparse_dense_shape)
#=============================================================================
# String Encoding
#=============================================================================
def __str__(self):
if self._is_eager():
- return "RaggedTensor(%s)" % self.tolist()
+ return "<tf.RaggedTensor %s>" % self.to_list()
else:
return self.__repr__()
def __repr__(self):
- return "RaggedTensor(values=%s, row_splits=%s)" % (self._values,
- self._row_splits)
+ return "tf.RaggedTensor(values=%s, row_splits=%s)" % (self._values,
+ self._row_splits)
#=============================================================================
# Eager Execution Mode
#=============================================================================
- def tolist(self):
+ def to_list(self):
"""Returns a nested Python `list` with the values for this `RaggedTensor`.
- If a `RaggedTensor` `rt` was constructed in graph execution mode, then
- `rt.tolist()` is equivalent to `rt.eval().tolist()`.
-
- If a `RaggedTensor` `rt` was constructed in eager execution mode, then
- `rt.tolist()` builds the Python list based on `rt`'s `EagerTensor`
- components.
+ Requires that `rt` was constructed in eager execution mode.
Returns:
A nested Python `list`.
"""
if self._is_eager():
- return self._eager_value().tolist()
+ return self._eager_value().to_list()
else:
- return self.eval().tolist()
+ raise ValueError("RaggedTensor.to_list() is only supported in eager "
+ "mode; in graph mode, evaluate the RaggedTensor first "
+ "and then use RaggedTensorValue.to_list().")
def _eager_value(self):
"""Returns a RaggedTensorValue for self. Requires self._is_eager()=true."""
- value = self.inner_values.numpy()
+ value = self.flat_values.numpy()
for row_splits in reversed(self.nested_row_splits):
value = ragged_tensor_value.RaggedTensorValue(value, row_splits.numpy())
return value
@@ -563,24 +1412,6 @@
return isinstance(rt, ops.EagerTensor)
#=============================================================================
- # Evaluation
- #=============================================================================
- def eval(self, feed_dict=None, session=None): # pylint: disable=redefined-outer-name
- """Evaluates this ragged tensor in a `Session`.
-
- Args:
- feed_dict: A dictionary that maps `Tensor` objects to feed values. See
- `tf.Session.run` for a description of the valid feed values.
- session: The `Session` to be used to evaluate this ragged tensor. If none,
- the default session will be used.
-
- Returns:
- A `RaggedTensorValue` object.
- """
- return _eval_using_default_session(self, feed_dict,
- self._as_graph_element().graph, session)
-
- #=============================================================================
# Indexing & Slicing
#=============================================================================
def __getitem__(self, key):
@@ -614,6 +1445,53 @@
#===============================================================================
+# Convert value -> tensor
+#===============================================================================
+def convert_to_tensor_or_ragged_tensor(value,
+ dtype=None,
+ preferred_dtype=None,
+ name=None):
+ """Converts value to a `RaggedTensor` or `Tensor`.
+
+ * If `value` is a `RaggedTensor`, then return it as-is.
+ * If `value` is a `RaggedTensorValue`, return a corresponding constant
+ `RaggedTensor`.
+ * Otherwise, use `convert_to_tensor` to convert `value` to a `Tensor`.
+
+ Args:
+ value: A `RaggedTensor`, a `RaggedTensorValue`, or an object whose type has
+ a registered `Tensor` conversion function.
+ dtype: Optional element type for the returned tensor. If missing the type
+ is inferred from the type of `value`.
+ preferred_dtype: Optional element type for the returned tensor, used when
+ dtype is None. This argument has no effect if `value` is already a
+ tensor, or when conversion is not possible.
+ name: Optional name to use if a new `Tensor` is created.
+
+ Returns:
+ A `Tensor` or `RaggedTensor`.
+ """
+ if isinstance(value, RaggedTensor):
+ if dtype and not dtype.is_compatible_with(value.dtype):
+ raise ValueError("Tensor conversion requested dtype %s for "
+ "RaggedTensor with dtype %s: %r" %
+ (dtype.name, value.dtype.name, value))
+ return value
+ elif isinstance(value, ragged_tensor_value.RaggedTensorValue):
+ with ops.name_scope(name, "ConvertToTensorOrRaggedTensor", []):
+ flat_values = ops.convert_to_tensor(
+ value=value.flat_values,
+ dtype=dtype,
+ preferred_dtype=preferred_dtype,
+ name="flat_values")
+ return RaggedTensor.from_nested_row_splits(flat_values,
+ value.nested_row_splits)
+ else:
+ return ops.convert_to_tensor(
+ value=value, dtype=dtype, preferred_dtype=preferred_dtype, name=name)
+
+
+#===============================================================================
# Register RaggedTensor for use with session.run.
#===============================================================================
def _ragged_tensor_value_from_components(components):
@@ -625,18 +1503,18 @@
def _ragged_tensor_session_fetch(rt):
- components = rt.nested_row_splits + (rt.inner_values,)
+ components = rt.nested_row_splits + (rt.flat_values,)
return (components, _ragged_tensor_value_from_components)
def _ragged_tensor_session_feed(feed_key, feed_val):
- key_components = feed_key.nested_row_splits + (feed_key.inner_values,)
- val_components = feed_val.nested_row_splits + (feed_val.inner_values,)
+ key_components = feed_key.nested_row_splits + (feed_key.flat_values,)
+ val_components = feed_val.nested_row_splits + (feed_val.flat_values,)
return zip(key_components, val_components)
def _ragged_tensor_session_feed_for_partial_run(feed_key):
- return feed_key.nested_row_splits + (feed_key.inner_values,)
+ return feed_key.nested_row_splits + (feed_key.flat_values,)
session.register_session_run_conversion_functions(
@@ -644,6 +1522,9 @@
_ragged_tensor_session_feed_for_partial_run)
+#===============================================================================
+# RaggedTensorType
+#===============================================================================
class RaggedTensorType(object):
"""Encoding of a static type for a `RaggedTensor`.
@@ -663,3 +1544,67 @@
dtype = property(lambda self: self._dtype)
ragged_rank = property(lambda self: self._ragged_rank)
+
+
+#===============================================================================
+# Helper Functions
+#===============================================================================
+def _assert_sparse_indices_are_ragged_right(indices):
+ """Checks that the given SparseTensor.indices tensor is ragged-right.
+
+ Example: `indices = [[0, 0], [0, 1], [2, 0], [3, 1]]` is not ragged right
+ because the entry `[3, 1]` skips a cell.
+
+ Args:
+ indices: The SparseTensor indices to check.
+
+ Returns:
+ A list of control dependency op tensors.
+ """
+ index_prefix = indices[:, :-1]
+ index_suffix = indices[:, -1]
+
+ # Check whether each index is starting a new row in the innermost dimension
+ # (prefix[i] != prefix[i-1]) or continuing a row (prefix[i] == prefix[i-1]).
+ # (Note: this skips the first index; we will check that separately below.)
+ index_prefix_changed = math_ops.reduce_any(
+ math_ops.not_equal(index_prefix[1:], index_prefix[:-1]), axis=1)
+
+ # Check two cases:
+ # * For indices that start a new row: index_suffix[i] must be zero.
+ # * For indices that continue a row: index_suffix[i] must be equal to
+ # index_suffix[i-1]+1.
+ index_ok = array_ops.where(
+ index_prefix_changed, math_ops.equal(index_suffix[1:], 0),
+ math_ops.equal(index_suffix[1:], index_suffix[:-1] + 1))
+
+ # Also check that the very first index didn't skip any cells. The first
+ # index starts a new row (by definition), so its suffix should be zero.
+ sparse_indices_are_ragged_right = math_ops.logical_and(
+ math_ops.reduce_all(math_ops.equal(index_suffix[:1], 0)),
+ math_ops.reduce_all(index_ok))
+
+ message = [
+ "SparseTensor is not right-ragged", "SparseTensor.indices =", indices
+ ]
+ return [control_flow_ops.Assert(sparse_indices_are_ragged_right, message)]
+
+
+@ops.RegisterGradient("RaggedTensorToSparse")
+def _ragged_tensor_to_sparse_gradient(op, unused_sparse_indices_grad,
+ sparse_values_grad,
+ unused_sparse_shape_grad):
+ """Gradient for RaggedTensorToSparse."""
+ op_inputs_nested_row_splits = op.inputs[:-1]
+ op_inputs_flat_values = op.inputs[-1]
+
+ # No gradient for the RaggedTensor's nested_row_splits.
+ nested_row_splits_gradient = [None] * len(op_inputs_nested_row_splits)
+
+ # Gradient for the RaggedTensor's flat_values is formed by reshaping
+ # the gradient for the SparseTensor's values.
+ flat_values_shape = array_ops.shape(op_inputs_flat_values)
+ flat_values_gradient = array_ops.reshape(sparse_values_grad,
+ flat_values_shape)
+
+ return nested_row_splits_gradient + [flat_values_gradient]
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_bounding_shape_op_test.py b/tensorflow/python/ops/ragged/ragged_tensor_bounding_shape_op_test.py
index befe30f..025a221 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_bounding_shape_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_bounding_shape_op_test.py
@@ -19,48 +19,45 @@
from __future__ import print_function
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedTensorBoundingShapeOp(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorBoundingShapeOp(ragged_test_util.RaggedTensorTestCase):
def testDocStringExample(self):
# This is the example from ragged.bounding_shape.__doc__.
- rt = ragged.constant([[1, 2, 3, 4], [5], [], [6, 7, 8, 9], [10]])
- self.assertEqual(self.evaluate(ragged.bounding_shape(rt)).tolist(), [5, 4])
+ rt = ragged_factory_ops.constant([[1, 2, 3, 4], [5], [], [6, 7, 8, 9],
+ [10]])
+ self.assertRaggedEqual(rt.bounding_shape(), [5, 4])
def test2DRaggedTensorWithOneRaggedDimension(self):
values = ['a', 'b', 'c', 'd', 'e', 'f', 'g']
- rt1 = ragged.from_row_splits(values, [0, 2, 5, 6, 6, 7])
- rt2 = ragged.from_row_splits(values, [0, 7])
- rt3 = ragged.from_row_splits(values, [0, 0, 7, 7])
- self.assertEqual(self.evaluate(ragged.bounding_shape(rt1)).tolist(), [5, 3])
- self.assertEqual(self.evaluate(ragged.bounding_shape(rt2)).tolist(), [1, 7])
- self.assertEqual(self.evaluate(ragged.bounding_shape(rt3)).tolist(), [3, 7])
+ rt1 = ragged_tensor.RaggedTensor.from_row_splits(values, [0, 2, 5, 6, 6, 7])
+ rt2 = ragged_tensor.RaggedTensor.from_row_splits(values, [0, 7])
+ rt3 = ragged_tensor.RaggedTensor.from_row_splits(values, [0, 0, 7, 7])
+ self.assertRaggedEqual(rt1.bounding_shape(), [5, 3])
+ self.assertRaggedEqual(rt2.bounding_shape(), [1, 7])
+ self.assertRaggedEqual(rt3.bounding_shape(), [3, 7])
def test3DRaggedTensorWithOneRaggedDimension(self):
values = [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]]
- rt1 = ragged.from_row_splits(values, [0, 2, 5, 6, 6, 7])
- rt2 = ragged.from_row_splits(values, [0, 7])
- rt3 = ragged.from_row_splits(values, [0, 0, 7, 7])
- self.assertEqual(
- self.evaluate(ragged.bounding_shape(rt1)).tolist(), [5, 3, 2])
- self.assertEqual(
- self.evaluate(ragged.bounding_shape(rt2)).tolist(), [1, 7, 2])
- self.assertEqual(
- self.evaluate(ragged.bounding_shape(rt3)).tolist(), [3, 7, 2])
-
- def testNonRaggedTensor(self):
- dt = [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
- self.assertEqual(self.evaluate(ragged.bounding_shape(dt)).tolist(), [4, 3])
+ rt1 = ragged_tensor.RaggedTensor.from_row_splits(values, [0, 2, 5, 6, 6, 7])
+ rt2 = ragged_tensor.RaggedTensor.from_row_splits(values, [0, 7])
+ rt3 = ragged_tensor.RaggedTensor.from_row_splits(values, [0, 0, 7, 7])
+ self.assertRaggedEqual(rt1.bounding_shape(), [5, 3, 2])
+ self.assertRaggedEqual(rt2.bounding_shape(), [1, 7, 2])
+ self.assertRaggedEqual(rt3.bounding_shape(), [3, 7, 2])
def testExplicitAxisOptimizations(self):
- rt = ragged.from_row_splits(b'a b c d e f g'.split(), [0, 2, 5, 6, 6, 7])
- self.assertEqual(self.evaluate(ragged.bounding_shape(rt, 0)).tolist(), 5)
- self.assertEqual(self.evaluate(ragged.bounding_shape(rt, 1)).tolist(), 3)
- self.assertEqual(
- self.evaluate(ragged.bounding_shape(rt, [1, 0])).tolist(), [3, 5])
+ rt = ragged_tensor.RaggedTensor.from_row_splits(b'a b c d e f g'.split(),
+ [0, 2, 5, 6, 6, 7])
+ self.assertRaggedEqual(rt.bounding_shape(0), 5)
+ self.assertRaggedEqual(rt.bounding_shape(1), 3)
+ self.assertRaggedEqual(rt.bounding_shape([1, 0]), [3, 5])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_shape.py b/tensorflow/python/ops/ragged/ragged_tensor_shape.py
index 9129b4b..706881d 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_shape.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_shape.py
@@ -21,13 +21,13 @@
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_util
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.ragged import ragged_array_ops
from tensorflow.python.ops.ragged import ragged_conversion_ops
-from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged import ragged_util
@@ -55,7 +55,7 @@
be ragged.
* "Inner dimensions" are dimensions that are encoded using a
- `RaggedTensor`'s `inner_values`. Inner dimensions are always uniform.
+ `RaggedTensor`'s `flat_values`. Inner dimensions are always uniform.
The sizes of partitioned dimensions are recorded using `partitioned_dim_sizes`
and `inner_dim_sizes`:
@@ -161,15 +161,15 @@
def from_tensor(cls, rt_input):
"""Constructs a ragged shape for a potentially ragged tensor."""
with ops.name_scope(None, 'RaggedTensorDynamicShapeFromTensor', [rt_input]):
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(rt_input)
+ rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt_input)
if not ragged_tensor.is_ragged(rt_input):
return cls([], array_ops.shape(rt_input))
else:
- partitioned_dim_sizes = ((ragged_array_ops.nrows(rt_input),) +
- ragged_array_ops.nested_row_lengths(rt_input))
+ partitioned_dim_sizes = (
+ (rt_input.nrows(),) + rt_input.nested_row_lengths())
return RaggedTensorDynamicShape(
partitioned_dim_sizes,
- array_ops.shape(rt_input.inner_values)[1:])
+ array_ops.shape(rt_input.flat_values)[1:])
def dimension_size(self, axis):
"""Returns the size of slices across the specified dimension."""
@@ -197,7 +197,7 @@
@property
def rank(self):
"""The number of dimensions in this shape, or None if unknown."""
- inner_ndims = self._inner_dim_sizes.shape[0].value
+ inner_ndims = tensor_shape.dimension_value(self._inner_dim_sizes.shape[0])
if inner_ndims is None:
return None
else:
@@ -229,7 +229,7 @@
@property
def num_inner_dimensions(self):
"""The number of inner dimensions, or `None` if not statically known."""
- return self._inner_dim_sizes.shape[0].value
+ return tensor_shape.dimension_value(self._inner_dim_sizes.shape[0])
def broadcast_to_rank(self, rank):
"""Adds leading size-1 dimensions to broadcast `self` to the given rank.
@@ -456,7 +456,7 @@
"""
if not isinstance(shape, RaggedTensorDynamicShape):
raise TypeError('shape must be a RaggedTensorDynamicShape')
- rt_input = ragged_factory_ops.convert_to_tensor_or_ragged_tensor(rt_input)
+ rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt_input)
# Broadcasting to a uniform shape.
if shape.num_partitioned_dimensions == 0:
@@ -497,17 +497,20 @@
rt_input = array_ops.reshape(
rt_input, array_ops.concat([[-1], dst_shape.inner_dim_sizes], axis=0))
for _ in range(dst_shape.rank - rt_input.shape.ndims):
- rt_input = ragged_factory_ops.from_row_lengths(
- rt_input, [ragged_array_ops.nrows(rt_input)])
+ if ragged_tensor.is_ragged(rt_input):
+ nrows = rt_input.nrows()
+ else:
+ nrows = array_ops.shape(rt_input, out_type=dtypes.int64)[0]
+ rt_input = ragged_tensor.RaggedTensor.from_row_lengths(rt_input, [nrows])
# Add ragged dimensions to match dst_shape.
if ragged_tensor.is_ragged(rt_input):
inner_rank_diff = (
- rt_input.inner_values.shape.ndims - 1 - dst_shape.num_inner_dimensions)
+ rt_input.flat_values.shape.ndims - 1 - dst_shape.num_inner_dimensions)
if inner_rank_diff > 0:
- rt_input = rt_input.with_inner_values(
+ rt_input = rt_input.with_flat_values(
ragged_conversion_ops.from_tensor(
- rt_input.inner_values, ragged_rank=inner_rank_diff))
+ rt_input.flat_values, ragged_rank=inner_rank_diff))
else:
rt_input = ragged_conversion_ops.from_tensor(
rt_input, ragged_rank=dst_shape.num_partitioned_dimensions - 1)
@@ -528,9 +531,9 @@
rt_input = ragged_array_ops.tile(rt_input, multiples)
if broadcast_inner_dimensions:
- rt_input = rt_input.with_inner_values(
+ rt_input = rt_input.with_flat_values(
array_ops.reshape(
- rt_input.inner_values,
+ rt_input.flat_values,
array_ops.concat([[-1], dst_shape.inner_dim_sizes], axis=0)))
# Do broadcasting for dimensions that become ragged. We must do these from
@@ -555,7 +558,7 @@
_ragged_tile_axis(rt_input.values, axis - 1, repeats))
else:
src_row_splits = rt_input.nested_row_splits
- src_row_lengths = ragged_array_ops.nested_row_lengths(rt_input)
+ src_row_lengths = rt_input.nested_row_lengths()
splits = src_row_splits[0]
dst_row_lengths = [repeats]
@@ -563,8 +566,7 @@
dst_row_lengths.append(
ragged_util.repeat_ranges(src_row_lengths[i], splits, repeats))
splits = array_ops.gather(src_row_splits[i], splits)
- dst_values = ragged_util.repeat_ranges(rt_input.inner_values, splits,
+ dst_values = ragged_util.repeat_ranges(rt_input.flat_values, splits,
repeats)
- return ragged_factory_ops.from_nested_row_lengths(dst_values,
- dst_row_lengths)
-
+ return ragged_tensor.RaggedTensor.from_nested_row_lengths(
+ dst_values, dst_row_lengths)
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_shape_test.py b/tensorflow/python/ops/ragged/ragged_tensor_shape_test.py
index 9c2dd26..bc0139c 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_shape_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_shape_test.py
@@ -23,27 +23,31 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.ops.ragged.ragged_tensor_shape import RaggedTensorDynamicShape
from tensorflow.python.platform import googletest
-class RaggedTensorShapeTest(test_util.TensorFlowTestCase,
- parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorBoundingShapeOp(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
def assertShapeEq(self, x, y):
- assert isinstance(x, ragged.RaggedTensorDynamicShape)
- assert isinstance(y, ragged.RaggedTensorDynamicShape)
+ assert isinstance(x, RaggedTensorDynamicShape)
+ assert isinstance(y, RaggedTensorDynamicShape)
x_partitioned_dim_sizes = [
- splits.eval().tolist() #
+ self.eval_to_list(splits) #
for splits in x.partitioned_dim_sizes
]
y_partitioned_dim_sizes = [
- splits.eval().tolist() #
+ self.eval_to_list(splits) #
for splits in y.partitioned_dim_sizes
]
self.assertEqual(x_partitioned_dim_sizes, y_partitioned_dim_sizes)
- self.assertEqual(x.inner_dim_sizes.eval().tolist(),
- y.inner_dim_sizes.eval().tolist())
+ self.assertAllEqual(x.inner_dim_sizes, y.inner_dim_sizes)
@parameterized.parameters([
dict(value='x', expected_dim_sizes=[]),
@@ -53,41 +57,41 @@
value=[[['a', 'b', 'c'], ['d', 'e', 'f']]],
expected_dim_sizes=[1, 2, 3]),
dict(
- value=ragged.constant_value([['a', 'b', 'c'], ['d', 'e']]),
+ value=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d',
+ 'e']]),
expected_dim_sizes=[2, [3, 2]]),
dict(
- value=ragged.constant_value([[['a', 'b', 'c'], ['d', 'e']]]),
+ value=ragged_factory_ops.constant_value([[['a', 'b', 'c'], ['d',
+ 'e']]]),
expected_dim_sizes=[1, [2], [3, 2]]),
dict(
- value=ragged.constant_value([[['a', 'b', 'c'], ['d', 'e', 'f']]],
- ragged_rank=1),
+ value=ragged_factory_ops.constant_value(
+ [[['a', 'b', 'c'], ['d', 'e', 'f']]], ragged_rank=1),
expected_dim_sizes=[1, [2], 3]),
dict(
- value=ragged.constant_value([[[[1], [2]], [[3], [4]]],
- [[[5], [6]]]], ragged_rank=1),
+ value=ragged_factory_ops.constant_value(
+ [[[[1], [2]], [[3], [4]]], [[[5], [6]]]], ragged_rank=1),
expected_dim_sizes=[2, [2, 1], 2, 1]),
dict(
- value=ragged.constant_value([[10, 20], [30]]),
+ value=ragged_factory_ops.constant_value([[10, 20], [30]]),
expected_dim_sizes=[2, [2, 1]]),
# Docstring examples:
dict(value=[[1, 2, 3], [4, 5, 6]], expected_dim_sizes=[2, 3]),
dict(
- value=ragged.constant_value([[1, 2], [], [3, 4, 5]]),
+ value=ragged_factory_ops.constant_value([[1, 2], [], [3, 4, 5]]),
expected_dim_sizes=[3, [2, 0, 3]]),
dict(
- value=ragged.constant_value([[[1, 2], [3, 4]], [[5, 6]]],
- ragged_rank=1),
+ value=ragged_factory_ops.constant_value([[[1, 2], [3, 4]], [[5, 6]]],
+ ragged_rank=1),
expected_dim_sizes=[2, [2, 1], 2]),
dict(
- value=ragged.constant_value([[[1, 2], [3]], [[4, 5]]]),
+ value=ragged_factory_ops.constant_value([[[1, 2], [3]], [[4, 5]]]),
expected_dim_sizes=[2, [2, 1], [2, 1, 2]]),
])
def testFromTensor(self, value, expected_dim_sizes):
- shape = ragged.RaggedTensorDynamicShape.from_tensor(value)
- expected = ragged.RaggedTensorDynamicShape.from_dim_sizes(
- expected_dim_sizes)
- with self.cached_session():
- self.assertShapeEq(shape, expected)
+ shape = RaggedTensorDynamicShape.from_tensor(value)
+ expected = RaggedTensorDynamicShape.from_dim_sizes(expected_dim_sizes)
+ self.assertShapeEq(shape, expected)
@parameterized.parameters([
dict(dim_sizes=[], rank=0, expected_dim_sizes=[]),
@@ -106,13 +110,11 @@
expected_dim_sizes=[1, 3, [3, 2, 4], 2, 3]),
])
def testBroadcastToRank(self, dim_sizes, rank, expected_dim_sizes):
- shape = ragged.RaggedTensorDynamicShape.from_dim_sizes(dim_sizes)
- expected = ragged.RaggedTensorDynamicShape.from_dim_sizes(
- expected_dim_sizes)
+ shape = RaggedTensorDynamicShape.from_dim_sizes(dim_sizes)
+ expected = RaggedTensorDynamicShape.from_dim_sizes(expected_dim_sizes)
broadcasted_shape = shape.broadcast_to_rank(rank)
- with self.cached_session():
- self.assertShapeEq(broadcasted_shape, expected)
- self.assertEqual(broadcasted_shape.rank, rank)
+ self.assertShapeEq(broadcasted_shape, expected)
+ self.assertEqual(broadcasted_shape.rank, rank)
@parameterized.parameters([
#=========================================================================
@@ -298,22 +300,19 @@
original_dim_sizes[axis] should be equal to `1` or `row_length`.
broadcast_dim_sizes: THe dimension sizes after broadcasting.
"""
- original_shape = ragged.RaggedTensorDynamicShape.from_dim_sizes(
- original_dim_sizes)
- broadcast_shape = ragged.RaggedTensorDynamicShape.from_dim_sizes(
- broadcast_dim_sizes)
- self.assertEqual(original_shape.rank, broadcast_shape.rank)
- with self.cached_session():
- # shape[axis].value == 1 and row_length > 1:
- bcast1 = original_shape.broadcast_dimension(axis, row_length)
- # shape[axis].value > 1 and row_length == shape[axis].value:
- bcast2 = broadcast_shape.broadcast_dimension(axis, row_length)
- # shape[axis].value > 1 and row_length == 1:
- bcast3 = broadcast_shape.broadcast_dimension(axis, 1)
+ original_shape = RaggedTensorDynamicShape.from_dim_sizes(original_dim_sizes)
+ bcast_shape = RaggedTensorDynamicShape.from_dim_sizes(broadcast_dim_sizes)
+ self.assertEqual(original_shape.rank, bcast_shape.rank)
+ # shape[axis].value == 1 and row_length > 1:
+ bcast1 = original_shape.broadcast_dimension(axis, row_length)
+ # shape[axis].value > 1 and row_length == shape[axis].value:
+ bcast2 = bcast_shape.broadcast_dimension(axis, row_length)
+ # shape[axis].value > 1 and row_length == 1:
+ bcast3 = bcast_shape.broadcast_dimension(axis, 1)
- self.assertShapeEq(bcast1, broadcast_shape)
- self.assertShapeEq(bcast2, broadcast_shape)
- self.assertShapeEq(bcast3, broadcast_shape)
+ self.assertShapeEq(bcast1, bcast_shape)
+ self.assertShapeEq(bcast2, bcast_shape)
+ self.assertShapeEq(bcast3, bcast_shape)
@parameterized.parameters(
[
@@ -371,116 +370,121 @@
expected_dims=[2, (2, 1), 2, (2, 1, 2, 1, 2, 1)]),
])
def testBroadcastDynamicShape(self, x_dims, y_dims, expected_dims):
- x_shape = ragged.RaggedTensorDynamicShape.from_dim_sizes(x_dims)
- y_shape = ragged.RaggedTensorDynamicShape.from_dim_sizes(y_dims)
- expected = ragged.RaggedTensorDynamicShape.from_dim_sizes(expected_dims)
- result1 = ragged.broadcast_dynamic_shape(x_shape, y_shape)
- result2 = ragged.broadcast_dynamic_shape(y_shape, x_shape)
- with self.cached_session():
- self.assertShapeEq(expected, result1)
- self.assertShapeEq(expected, result2)
+ x_shape = RaggedTensorDynamicShape.from_dim_sizes(x_dims)
+ y_shape = RaggedTensorDynamicShape.from_dim_sizes(y_dims)
+ expected = RaggedTensorDynamicShape.from_dim_sizes(expected_dims)
+ result1 = ragged_tensor_shape.broadcast_dynamic_shape(x_shape, y_shape)
+ result2 = ragged_tensor_shape.broadcast_dynamic_shape(y_shape, x_shape)
+ self.assertShapeEq(expected, result1)
+ self.assertShapeEq(expected, result2)
def testRepr(self):
- shape = ragged.RaggedTensorDynamicShape.from_dim_sizes([2, (2, 1), 2, 1])
+ shape = RaggedTensorDynamicShape.from_dim_sizes([2, (2, 1), 2, 1])
self.assertRegexpMatches(
repr(shape),
r'RaggedTensorDynamicShape\('
r'partitioned_dim_sizes=\(<[^>]+>, <[^>]+>\), '
r'inner_dim_sizes=<[^>]+>\)')
- @parameterized.parameters([
- dict(
- x=[[10], [20], [30]], # shape=[3, 1]
- dim_sizes=[3, 2],
- expected=[[10, 10], [20, 20], [30, 30]]),
- dict(
- x=[[10], [20], [30]], # shape=[3, 1]
- dim_sizes=[3, [3, 0, 2]],
- expected=ragged.constant_value([[10, 10, 10], [], [30, 30]],
- dtype=np.int32)),
- dict(
- x=[[[1, 2, 3]], [[4, 5, 6]]], # shape = [2, 1, 3]
- dim_sizes=[2, [2, 3], 3],
- expected=ragged.constant_value(
- [[[1, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6], [4, 5, 6]]],
- dtype=np.int32,
- ragged_rank=1)),
- dict(
- x=[[[1]], [[2]]], # shape = [2, 1, 1]
- dim_sizes=[2, [2, 3], [0, 2, 1, 2, 0]],
- expected=ragged.constant_value([[[], [1, 1]], [[2], [2, 2], []]],
- dtype=np.int32,
- ragged_rank=2)),
- dict(
- x=10,
- dim_sizes=[3, [3, 0, 2]],
- expected=ragged.constant_value([[10, 10, 10], [], [10, 10]])),
- ])
+ @parameterized.parameters(
+ [
+ dict(
+ x=[[10], [20], [30]], # shape=[3, 1]
+ dim_sizes=[3, 2],
+ expected=[[10, 10], [20, 20], [30, 30]]),
+ dict(
+ x=[[10], [20], [30]], # shape=[3, 1]
+ dim_sizes=[3, [3, 0, 2]],
+ expected=ragged_factory_ops.constant_value(
+ [[10, 10, 10], [], [30, 30]], dtype=np.int32)),
+ dict(
+ x=[[[1, 2, 3]], [[4, 5, 6]]], # shape = [2, 1, 3]
+ dim_sizes=[2, [2, 3], 3],
+ expected=ragged_factory_ops.constant_value(
+ [[[1, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6], [4, 5, 6]]],
+ dtype=np.int32,
+ ragged_rank=1)),
+ dict(
+ x=[[[1]], [[2]]], # shape = [2, 1, 1]
+ dim_sizes=[2, [2, 3], [0, 2, 1, 2, 0]],
+ expected=ragged_factory_ops.constant_value(
+ [[[], [1, 1]], [[2], [2, 2], []]],
+ dtype=np.int32,
+ ragged_rank=2)),
+ dict(
+ x=10,
+ dim_sizes=[3, [3, 0, 2]],
+ expected=ragged_factory_ops.constant_value([[10, 10, 10], [],
+ [10, 10]])),
+ ])
def testRaggedBroadcastTo(self, x, dim_sizes, expected):
- shape = ragged.RaggedTensorDynamicShape.from_dim_sizes(dim_sizes)
- result = ragged.broadcast_to(x, shape)
- with self.cached_session():
- self.assertEqual(
- getattr(result, 'ragged_rank', 0), getattr(expected, 'ragged_rank',
- 0))
- if hasattr(expected, 'tolist'):
- expected = expected.tolist()
- self.assertEqual(result.eval().tolist(), expected)
+ shape = RaggedTensorDynamicShape.from_dim_sizes(dim_sizes)
+ result = ragged_tensor_shape.broadcast_to(x, shape)
+ self.assertEqual(
+ getattr(result, 'ragged_rank', 0), getattr(expected, 'ragged_rank', 0))
+ self.assertRaggedEqual(result, expected)
- @parameterized.parameters([
- dict(
- doc='x.shape=[3, (D1)]; y.shape=[3, 1]; bcast.shape=[3, (D1)]',
- x=ragged.constant_value([[1, 2, 3], [], [4, 5]], dtype=np.int32),
- y=[[10], [20], [30]],
- expected=ragged.constant_value([[11, 12, 13], [], [34, 35]])),
- dict(
- doc='x.shape=[3, (D1)]; y.shape=[]; bcast.shape=[3, (D1)]',
- x=ragged.constant_value([[1, 2, 3], [], [4, 5]], dtype=np.int32),
- y=10,
- expected=ragged.constant_value([[11, 12, 13], [], [14, 15]])),
- dict(
- doc='x.shape=[1, (D1)]; y.shape=[3, 1]; bcast.shape=[3, (D1)]',
- x=ragged.constant_value([[1, 2, 3]], dtype=np.int32),
- y=[[10], [20], [30]],
- expected=ragged.constant_value(
- [[11, 12, 13], [21, 22, 23], [31, 32, 33]], dtype=np.int32)),
- dict(
- doc=('x.shape=[2, (D1), 1]; y.shape=[1, (D2)]; '
- 'bcast.shape=[2, (D1), (D2)]'),
- x=ragged.constant_value([[[1], [2], [3]], [[4]]], ragged_rank=1),
- y=ragged.constant_value([[10, 20, 30]]),
- expected=ragged.constant_value([[[11, 21, 31], [12, 22, 32],
- [13, 23, 33]], [[14, 24, 34]]])),
- dict(
- doc=('x.shape=[2, (D1), 1]; y.shape=[1, 1, 4]; '
- 'bcast.shape=[2, (D1), 4]'),
- x=ragged.constant_value([[[10], [20]], [[30]]], ragged_rank=1),
- y=[[[1, 2, 3, 4]]],
- expected=ragged.constant_value(
- [[[11, 12, 13, 14], [21, 22, 23, 24]], [[31, 32, 33, 34]]],
- ragged_rank=1)),
- dict(
- doc=('x.shape=[2, (D1), 2, 1]; y.shape=[2, (D2)]; '
- 'bcast.shape=[2, (D1), (2), (D2)'),
- x=ragged.constant_value([[[[1], [2]], [[3], [4]]],
- [[[5], [6]]]],
- ragged_rank=1),
- y=ragged.constant_value([[10, 20], [30]]),
- expected=ragged.constant_value(
- [[[[11, 21], [32]], [[13, 23], [34]]],
- [[[15, 25], [36]]]])),
- ])
+ @parameterized.parameters(
+ [
+ dict(
+ doc='x.shape=[3, (D1)]; y.shape=[3, 1]; bcast.shape=[3, (D1)]',
+ x=ragged_factory_ops.constant_value([[1, 2, 3], [], [4, 5]],
+ dtype=np.int32),
+ y=[[10], [20], [30]],
+ expected=ragged_factory_ops.constant_value([[11, 12, 13], [],
+ [34, 35]])),
+ dict(
+ doc='x.shape=[3, (D1)]; y.shape=[]; bcast.shape=[3, (D1)]',
+ x=ragged_factory_ops.constant_value([[1, 2, 3], [], [4, 5]],
+ dtype=np.int32),
+ y=10,
+ expected=ragged_factory_ops.constant_value([[11, 12, 13], [],
+ [14, 15]])),
+ dict(
+ doc='x.shape=[1, (D1)]; y.shape=[3, 1]; bcast.shape=[3, (D1)]',
+ x=ragged_factory_ops.constant_value([[1, 2, 3]], dtype=np.int32),
+ y=[[10], [20], [30]],
+ expected=ragged_factory_ops.constant_value(
+ [[11, 12, 13], [21, 22, 23], [31, 32, 33]], dtype=np.int32)),
+ dict(
+ doc=('x.shape=[2, (D1), 1]; y.shape=[1, (D2)]; '
+ 'bcast.shape=[2, (D1), (D2)]'),
+ x=ragged_factory_ops.constant_value([[[1], [2], [3]], [[4]]],
+ ragged_rank=1),
+ y=ragged_factory_ops.constant_value([[10, 20, 30]]),
+ expected=ragged_factory_ops.constant_value([[[11, 21, 31],
+ [12, 22, 32],
+ [13, 23, 33]],
+ [[14, 24, 34]]])),
+ dict(
+ doc=('x.shape=[2, (D1), 1]; y.shape=[1, 1, 4]; '
+ 'bcast.shape=[2, (D1), 4]'),
+ x=ragged_factory_ops.constant_value([[[10], [20]], [[30]]],
+ ragged_rank=1),
+ y=[[[1, 2, 3, 4]]],
+ expected=ragged_factory_ops.constant_value(
+ [[[11, 12, 13, 14], [21, 22, 23, 24]], [[31, 32, 33, 34]]],
+ ragged_rank=1)),
+ dict(
+ doc=('x.shape=[2, (D1), 2, 1]; y.shape=[2, (D2)]; '
+ 'bcast.shape=[2, (D1), (2), (D2)'),
+ x=ragged_factory_ops.constant_value(
+ [[[[1], [2]], [[3], [4]]], [[[5], [6]]]], ragged_rank=1),
+ y=ragged_factory_ops.constant_value([[10, 20], [30]]),
+ expected=ragged_factory_ops.constant_value([[[[11, 21], [32]],
+ [[13, 23], [34]]],
+ [[[15, 25], [36]]]])),
+ ])
def testRaggedAddWithBroadcasting(self, x, y, expected, doc):
expected_rrank = getattr(expected, 'ragged_rank', 0)
- x = ragged.convert_to_tensor_or_ragged_tensor(x, dtype=dtypes.int32)
- y = ragged.convert_to_tensor_or_ragged_tensor(y, dtype=dtypes.int32)
+ x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x, dtype=dtypes.int32)
+ y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y, dtype=dtypes.int32)
result = x + y
result_rrank = getattr(result, 'ragged_rank', 0)
self.assertEqual(expected_rrank, result_rrank)
if hasattr(expected, 'tolist'):
expected = expected.tolist()
- with self.cached_session():
- self.assertEqual(result.eval().tolist(), expected)
+ self.assertRaggedEqual(result, expected)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_test.py b/tensorflow/python/ops/ragged/ragged_tensor_test.py
index 608fbd6..89691b0 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_test.py
@@ -19,17 +19,22 @@
from __future__ import print_function
import re
-import sys
from absl.testing import parameterized
import numpy as np
+from tensorflow.python.eager import context
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 test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.ops.ragged import ragged_test_util
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
from tensorflow.python.platform import googletest
@@ -103,63 +108,62 @@
EXAMPLE_RAGGED_TENSOR_4D_SPLITS1 = [0, 2, 2, 3, 4]
EXAMPLE_RAGGED_TENSOR_4D_SPLITS2 = [0, 3, 6, 9, 10]
EXAMPLE_RAGGED_TENSOR_4D_VALUES = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10],
- [11, 12], [13, 14], [15, 16], [17,
- 18], [19, 20]]
+ [11, 12], [13, 14], [15, 16], [17, 18],
+ [19, 20]]
-class RaggedTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
longMessage = True # Property in unittest.Testcase. pylint: disable=invalid-name
#=============================================================================
# RaggedTensor class docstring examples
#=============================================================================
- @test_util.run_deprecated_v1
def testClassDocStringExamples(self):
# From section: "Component Tensors"
- rt = ragged.from_row_splits(
+ rt = RaggedTensor.from_row_splits(
values=[3, 1, 4, 1, 5, 9, 2, 6], row_splits=[0, 4, 4, 7, 8, 8])
- self.assertEqual(
- self.evaluate(rt).tolist(), [[3, 1, 4, 1], [], [5, 9, 2], [6], []])
+ self.assertRaggedEqual(rt, [[3, 1, 4, 1], [], [5, 9, 2], [6], []])
del rt
# From section: "Alternative Row-Partitioning Schemes"
values = [3, 1, 4, 1, 5, 9, 2, 6]
- rt1 = ragged.from_row_splits(values, row_splits=[0, 4, 4, 7, 8, 8])
- rt2 = ragged.from_row_lengths(values, row_lengths=[4, 0, 3, 1, 0])
- rt3 = ragged.from_value_rowids(
+ rt1 = RaggedTensor.from_row_splits(values, row_splits=[0, 4, 4, 7, 8, 8])
+ rt2 = RaggedTensor.from_row_lengths(values, row_lengths=[4, 0, 3, 1, 0])
+ rt3 = RaggedTensor.from_value_rowids(
values, value_rowids=[0, 0, 0, 0, 2, 2, 2, 3], nrows=5)
- rt4 = ragged.from_row_starts(values, row_starts=[0, 4, 4, 7, 8])
- rt5 = ragged.from_row_limits(values, row_limits=[4, 4, 7, 8, 8])
+ rt4 = RaggedTensor.from_row_starts(values, row_starts=[0, 4, 4, 7, 8])
+ rt5 = RaggedTensor.from_row_limits(values, row_limits=[4, 4, 7, 8, 8])
for rt in (rt1, rt2, rt3, rt4, rt5):
- self.assertEqual(
- self.evaluate(rt).tolist(), [[3, 1, 4, 1], [], [5, 9, 2], [6], []])
+ self.assertRaggedEqual(rt, [[3, 1, 4, 1], [], [5, 9, 2], [6], []])
del rt1, rt2, rt3, rt4, rt5
# From section: "Multiple Ragged Dimensions"
- inner_rt = ragged.from_row_splits(
+ inner_rt = RaggedTensor.from_row_splits(
values=[3, 1, 4, 1, 5, 9, 2, 6], row_splits=[0, 4, 4, 7, 8, 8])
- outer_rt = ragged.from_row_splits(values=inner_rt, row_splits=[0, 3, 3, 5])
+ outer_rt = RaggedTensor.from_row_splits(
+ values=inner_rt, row_splits=[0, 3, 3, 5])
self.assertEqual(outer_rt.ragged_rank, 2)
self.assertEqual(
- self.evaluate(outer_rt).tolist(),
+ self.eval_to_list(outer_rt),
[[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]])
del inner_rt, outer_rt
# From section: "Multiple Ragged Dimensions"
- rt = ragged.from_nested_row_splits(
- inner_values=[3, 1, 4, 1, 5, 9, 2, 6],
+ rt = RaggedTensor.from_nested_row_splits(
+ flat_values=[3, 1, 4, 1, 5, 9, 2, 6],
nested_row_splits=([0, 3, 3, 5], [0, 4, 4, 7, 8, 8]))
self.assertEqual(
- self.evaluate(rt).tolist(),
- [[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]])
+ self.eval_to_list(rt), [[[3, 1, 4, 1], [], [5, 9, 2]], [], [[6], []]])
del rt
# From section: "Uniform Inner Dimensions"
- rt = ragged.from_row_splits(
+ rt = RaggedTensor.from_row_splits(
values=array_ops.ones([5, 3]), row_splits=[0, 2, 5])
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[[1, 1, 1], [1, 1, 1]], [[1, 1, 1], [1, 1, 1], [1, 1, 1]]])
self.assertEqual(rt.shape.as_list(), [2, None, 3])
del rt
@@ -174,18 +178,19 @@
splits2 = np.array([0, 3, 5], dtype=np.int64)
# Test construction of a RaggedTensorValue with ragged_rank=1.
- rt_value = ragged.RaggedTensorValue(values, splits)
+ rt_value = ragged_tensor_value.RaggedTensorValue(values, splits)
self.assertEqual(rt_value.row_splits.dtype, np.int64)
self.assertEqual(rt_value.shape, (5, None))
self.assertEqual(len(rt_value.nested_row_splits), 1)
self.assertAllEqual(splits, rt_value.row_splits)
self.assertAllEqual(values, rt_value.values)
self.assertAllEqual(splits, rt_value.nested_row_splits[0])
- self.assertAllEqual(values, rt_value.inner_values)
+ self.assertAllEqual(values, rt_value.flat_values)
# Test construction of a RaggedTensorValue with ragged_rank=2.
- rt_value = ragged.RaggedTensorValue(
- values=ragged.RaggedTensorValue(values, splits), row_splits=splits2)
+ rt_value = ragged_tensor_value.RaggedTensorValue(
+ values=ragged_tensor_value.RaggedTensorValue(values, splits),
+ row_splits=splits2)
self.assertEqual(rt_value.row_splits.dtype, np.int64)
self.assertEqual(rt_value.shape, (2, None, None))
self.assertEqual(len(rt_value.nested_row_splits), 2)
@@ -194,21 +199,19 @@
self.assertAllEqual(splits2, rt_value.nested_row_splits[0])
self.assertAllEqual(splits, rt_value.nested_row_splits[1])
self.assertAllEqual(values, rt_value.values.values)
- self.assertAllEqual(values, rt_value.inner_values)
+ self.assertAllEqual(values, rt_value.flat_values)
#=============================================================================
# RaggedTensor Constructor (private)
#=============================================================================
- @test_util.run_deprecated_v1
def testRaggedTensorConstruction(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
- rt = ragged.RaggedTensor(
- values=values, row_splits=row_splits, internal=True)
+ rt = RaggedTensor(values=values, row_splits=row_splits, internal=True)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
def testRaggedTensorConstructionErrors(self):
@@ -217,117 +220,118 @@
with self.assertRaisesRegexp(ValueError,
'RaggedTensor constructor is private'):
- ragged.RaggedTensor(values=values, row_splits=row_splits)
+ RaggedTensor(values=values, row_splits=row_splits)
with self.assertRaisesRegexp(TypeError,
'values must be a Tensor or RaggedTensor'):
- ragged.RaggedTensor(values=range(7), row_splits=row_splits, internal=True)
+ RaggedTensor(values=range(7), row_splits=row_splits, internal=True)
with self.assertRaisesRegexp(TypeError,
'Row-partitioning argument must be a Tensor'):
- ragged.RaggedTensor(
- values=values, row_splits=[0, 2, 2, 5, 6, 7], internal=True)
+ RaggedTensor(values=values, row_splits=[0, 2, 2, 5, 6, 7], internal=True)
with self.assertRaisesRegexp(ValueError,
r'Shape \(6, 1\) must have rank 1'):
- ragged.RaggedTensor(
+ RaggedTensor(
values=values,
row_splits=array_ops.expand_dims(row_splits, 1),
internal=True)
with self.assertRaisesRegexp(TypeError,
'Cached value must be a Tensor or None.'):
- ragged.RaggedTensor(values=values, row_splits=row_splits,
- cached_row_lengths=[2, 3, 4], internal=True)
+ RaggedTensor(
+ values=values,
+ row_splits=row_splits,
+ cached_row_lengths=[2, 3, 4],
+ internal=True)
#=============================================================================
# RaggedTensor Factory Ops
#=============================================================================
- @test_util.run_deprecated_v1
def testFromValueRowIdsWithDerivedNRows(self):
# nrows is known at graph creation time.
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
- rt = ragged.from_value_rowids(values, value_rowids)
+ rt = RaggedTensor.from_value_rowids(values, value_rowids)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [5, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_value_rowids = ragged.value_rowids(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_value_rowids = rt.value_rowids()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
self.assertIs(rt_value_rowids, value_rowids) # cached_value_rowids
self.assertAllEqual(rt_value_rowids, value_rowids)
- self.assertEqual(self.evaluate(rt_nrows), 5)
+ self.assertEqual(self.eval_to_list(rt_nrows), 5)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- @test_util.run_deprecated_v1
def testFromValueRowIdsWithDerivedNRowsDynamic(self):
# nrows is not known at graph creation time.
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
value_rowids = array_ops.placeholder_with_default(value_rowids, shape=None)
- rt = ragged.from_value_rowids(values, value_rowids)
+ rt = RaggedTensor.from_value_rowids(values, value_rowids)
self.assertEqual(rt.dtype, dtypes.string)
- self.assertEqual(rt.shape.as_list(), [None, None])
+ if context.executing_eagerly():
+ self.assertEqual(rt.shape.as_list(), [5, None])
+ else:
+ self.assertEqual(rt.shape.as_list(), [None, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_value_rowids = ragged.value_rowids(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_value_rowids = rt.value_rowids()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
self.assertIs(rt_value_rowids, value_rowids) # cached_value_rowids
self.assertAllEqual(rt_value_rowids, value_rowids)
- self.assertEqual(self.evaluate(rt_nrows), 5)
+ self.assertEqual(self.eval_to_list(rt_nrows), 5)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- @test_util.run_deprecated_v1
def testFromValueRowIdsWithExplicitNRows(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
nrows = constant_op.constant(7, dtypes.int64)
- rt = ragged.from_value_rowids(values, value_rowids, nrows)
+ rt = RaggedTensor.from_value_rowids(values, value_rowids, nrows)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [7, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_value_rowids = ragged.value_rowids(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_value_rowids = rt.value_rowids()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
self.assertIs(rt_value_rowids, value_rowids) # cached_value_rowids
self.assertIs(rt_nrows, nrows) # cached_nrows
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g'], [], []])
- @test_util.run_deprecated_v1
def testFromValueRowIdsWithExplicitNRowsEqualToDefault(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
nrows = constant_op.constant(5, dtypes.int64)
- rt = ragged.from_value_rowids(values, value_rowids, nrows)
+ rt = RaggedTensor.from_value_rowids(values, value_rowids, nrows)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [5, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_value_rowids = ragged.value_rowids(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_value_rowids = rt.value_rowids()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
self.assertIs(rt_value_rowids, value_rowids) # cached_value_rowids
@@ -335,112 +339,106 @@
self.assertAllEqual(rt_value_rowids, value_rowids)
self.assertAllEqual(rt_nrows, nrows)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- @test_util.run_deprecated_v1
def testFromValueRowIdsWithEmptyValues(self):
- rt = ragged.from_value_rowids([], [])
- rt_nrows = ragged.nrows(rt)
+ rt = RaggedTensor.from_value_rowids([], [])
+ rt_nrows = rt.nrows()
self.assertEqual(rt.dtype, dtypes.float32)
self.assertEqual(rt.shape.as_list(), [0, None])
self.assertEqual(rt.ragged_rank, 1)
self.assertEqual(rt.values.shape.as_list(), [0])
- self.assertEqual(ragged.value_rowids(rt).shape.as_list(), [0])
- self.assertEqual(self.evaluate(rt_nrows).tolist(), 0)
- self.assertEqual(self.evaluate(rt).tolist(), [])
+ self.assertEqual(rt.value_rowids().shape.as_list(), [0])
+ self.assertEqual(self.eval_to_list(rt_nrows), 0)
+ self.assertEqual(self.eval_to_list(rt), [])
- @test_util.run_deprecated_v1
def testFromRowSplits(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
- rt = ragged.from_row_splits(values, row_splits)
+ rt = RaggedTensor.from_row_splits(values, row_splits)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [5, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
rt_row_splits = rt.row_splits
- rt_nrows = ragged.nrows(rt)
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
self.assertIs(rt_row_splits, row_splits)
- self.assertEqual(self.evaluate(rt_nrows), 5)
+ self.assertEqual(self.eval_to_list(rt_nrows), 5)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
def testFromRowSplitsWithEmptySplits(self):
err_msg = 'row_splits tensor may not be empty'
with self.assertRaisesRegexp(ValueError, err_msg):
- ragged.from_row_splits([], [])
+ RaggedTensor.from_row_splits([], [])
- @test_util.run_deprecated_v1
def testFromRowStarts(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
row_starts = constant_op.constant([0, 2, 2, 5, 6], dtypes.int64)
- rt = ragged.from_row_starts(values, row_starts)
+ rt = RaggedTensor.from_row_starts(values, row_starts)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [5, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_row_starts = ragged.row_starts(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_row_starts = rt.row_starts()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
- self.assertEqual(self.evaluate(rt_nrows), 5)
+ self.assertEqual(self.eval_to_list(rt_nrows), 5)
self.assertAllEqual(rt_row_starts, row_starts)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- @test_util.run_deprecated_v1
def testFromRowLimits(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
row_limits = constant_op.constant([2, 2, 5, 6, 7], dtypes.int64)
- rt = ragged.from_row_limits(values, row_limits)
+ rt = RaggedTensor.from_row_limits(values, row_limits)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [5, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_row_limits = ragged.row_limits(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_row_limits = rt.row_limits()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
- self.assertEqual(self.evaluate(rt_nrows), 5)
+ self.assertEqual(self.eval_to_list(rt_nrows), 5)
self.assertAllEqual(rt_row_limits, row_limits)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- @test_util.run_deprecated_v1
def testFromRowLengths(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
row_lengths = constant_op.constant([2, 0, 3, 1, 1], dtypes.int64)
- rt = ragged.from_row_lengths(values, row_lengths)
+ rt = RaggedTensor.from_row_lengths(values, row_lengths)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [5, None])
self.assertEqual(rt.ragged_rank, 1)
rt_values = rt.values
- rt_row_lengths = ragged.row_lengths(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_row_lengths = rt.row_lengths()
+ rt_nrows = rt.nrows()
self.assertIs(rt_values, values)
self.assertIs(rt_row_lengths, row_lengths) # cached_nrows
- self.assertEqual(self.evaluate(rt_nrows), 5)
+ self.assertEqual(self.eval_to_list(rt_nrows), 5)
self.assertAllEqual(rt_row_lengths, row_lengths)
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- @test_util.run_deprecated_v1
def testFromNestedValueRowIdsWithDerivedNRows(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
nested_value_rowids = [
@@ -448,24 +446,23 @@
constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
]
- rt = ragged.from_nested_value_rowids(values, nested_value_rowids)
+ rt = RaggedTensor.from_nested_value_rowids(values, nested_value_rowids)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [4, None, None])
self.assertEqual(rt.ragged_rank, 2)
rt_values = rt.values
- rt_value_rowids = ragged.value_rowids(rt)
+ rt_value_rowids = rt.value_rowids()
rt_values_values = rt_values.values
- rt_values_value_rowids = ragged.value_rowids(rt_values)
+ rt_values_value_rowids = rt_values.value_rowids()
self.assertIs(rt_values_values, values)
self.assertAllEqual(rt_value_rowids, nested_value_rowids[0])
self.assertAllEqual(rt_values_value_rowids, nested_value_rowids[1])
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[[b'a', b'b'], []], [[b'c', b'd', b'e']], [], [[b'f'], [b'g']]])
- @test_util.run_deprecated_v1
def testFromNestedValueRowIdsWithExplicitNRows(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
nested_value_rowids = [
@@ -477,17 +474,18 @@
constant_op.constant(6, dtypes.int64)
]
- rt = ragged.from_nested_value_rowids(values, nested_value_rowids, nrows)
+ rt = RaggedTensor.from_nested_value_rowids(values, nested_value_rowids,
+ nrows)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [6, None, None])
self.assertEqual(rt.ragged_rank, 2)
rt_values = rt.values
- rt_value_rowids = ragged.value_rowids(rt)
- rt_nrows = ragged.nrows(rt)
+ rt_value_rowids = rt.value_rowids()
+ rt_nrows = rt.nrows()
rt_values_values = rt_values.values
- rt_values_value_rowids = ragged.value_rowids(rt_values)
- rt_values_nrows = ragged.nrows(rt_values)
+ rt_values_value_rowids = rt_values.value_rowids()
+ rt_values_nrows = rt_values.nrows()
self.assertIs(rt_values_values, values)
self.assertAllEqual(rt_value_rowids, nested_value_rowids[0])
@@ -495,9 +493,8 @@
self.assertAllEqual(rt_nrows, nrows[0])
self.assertAllEqual(rt_values_nrows, nrows[1])
self.assertEqual(
- self.evaluate(rt).tolist(),
- [[[b'a', b'b'], []], [[b'c', b'd', b'e']], [], [[b'f'], [b'g'], []], [],
- []])
+ self.eval_to_list(rt), [[[b'a', b'b'], []], [[b'c', b'd', b'e']], [],
+ [[b'f'], [b'g'], []], [], []])
def testFromNestedValueRowIdsWithExplicitNRowsMismatch(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
@@ -509,28 +506,26 @@
with self.assertRaisesRegexp(
ValueError, 'nested_nrows must have the same '
'length as nested_value_rowids'):
- ragged.from_nested_value_rowids(values, nested_value_rowids, nrows)
+ RaggedTensor.from_nested_value_rowids(values, nested_value_rowids, nrows)
def testFromNestedValueRowIdsWithNonListInput(self):
with self.assertRaisesRegexp(
TypeError, 'nested_value_rowids must be a list of Tensors'):
- ragged.from_nested_value_rowids([1, 2, 3],
- constant_op.constant(
- [[0, 1, 2], [0, 1, 2]], dtypes.int64))
+ RaggedTensor.from_nested_value_rowids(
+ [1, 2, 3], constant_op.constant([[0, 1, 2], [0, 1, 2]], dtypes.int64))
with self.assertRaisesRegexp(TypeError,
'nested_nrows must be a list of Tensors'):
- ragged.from_nested_value_rowids([1, 2, 3], [[0, 1, 2], [0, 1, 2]],
- constant_op.constant([3, 3]))
+ RaggedTensor.from_nested_value_rowids([1, 2, 3], [[0, 1, 2], [0, 1, 2]],
+ constant_op.constant([3, 3]))
- @test_util.run_deprecated_v1
def testFromNestedRowSplits(self):
- inner_values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
+ flat_values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
nested_row_splits = [
constant_op.constant([0, 2, 3, 3, 5], dtypes.int64),
constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
]
- rt = ragged.from_nested_row_splits(inner_values, nested_row_splits)
+ rt = RaggedTensor.from_nested_row_splits(flat_values, nested_row_splits)
self.assertEqual(rt.dtype, dtypes.string)
self.assertEqual(rt.shape.as_list(), [4, None, None])
self.assertEqual(rt.ragged_rank, 2)
@@ -540,19 +535,18 @@
rt_values_values = rt_values.values
rt_values_row_splits = rt_values.row_splits
- self.assertIs(rt_values_values, inner_values)
+ self.assertIs(rt_values_values, flat_values)
self.assertIs(rt_row_splits, nested_row_splits[0])
self.assertIs(rt_values_row_splits, nested_row_splits[1])
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[[b'a', b'b'], []], [[b'c', b'd', b'e']], [], [[b'f'], [b'g']]])
def testFromNestedRowSplitsWithNonListInput(self):
with self.assertRaisesRegexp(TypeError,
'nested_row_splits must be a list of Tensors'):
- ragged.from_nested_row_splits([1, 2],
- constant_op.constant([[0, 1, 2], [0, 1, 2]],
- dtypes.int64))
+ RaggedTensor.from_nested_row_splits(
+ [1, 2], constant_op.constant([[0, 1, 2], [0, 1, 2]], dtypes.int64))
def testFromValueRowIdsWithBadNRows(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
@@ -560,7 +554,7 @@
nrows = constant_op.constant(5, dtypes.int64)
with self.assertRaisesRegexp(ValueError, r'Expected nrows >= 0; got -2'):
- ragged.from_value_rowids(
+ RaggedTensor.from_value_rowids(
values=values,
value_rowids=array_ops.placeholder_with_default(value_rowids, None),
nrows=-2)
@@ -568,113 +562,94 @@
with self.assertRaisesRegexp(
ValueError, r'Expected nrows >= value_rowids\[-1\] \+ 1; got nrows=2, '
r'value_rowids\[-1\]=4'):
- ragged.from_value_rowids(
+ RaggedTensor.from_value_rowids(
values=values, value_rowids=value_rowids, nrows=2)
with self.assertRaisesRegexp(
ValueError, r'Expected nrows >= value_rowids\[-1\] \+ 1; got nrows=4, '
r'value_rowids\[-1\]=4'):
- ragged.from_value_rowids(
+ RaggedTensor.from_value_rowids(
values=values, value_rowids=value_rowids, nrows=4)
with self.assertRaisesRegexp(ValueError,
r'Shape \(7, 1\) must have rank 1'):
- ragged.from_value_rowids(
+ RaggedTensor.from_value_rowids(
values=values,
value_rowids=array_ops.expand_dims(value_rowids, 1),
nrows=nrows)
with self.assertRaisesRegexp(ValueError, r'Shape \(1,\) must have rank 0'):
- ragged.from_value_rowids(
+ RaggedTensor.from_value_rowids(
values=values,
value_rowids=value_rowids,
nrows=array_ops.expand_dims(nrows, 0))
- @test_util.run_deprecated_v1
def testGraphMismatch(self):
- with ops.Graph().as_default():
- values = constant_op.constant([1, 2, 3])
- with ops.Graph().as_default():
- splits = constant_op.constant([0, 2, 3])
- self.assertRaisesRegexp(ValueError, '.* must be from the same graph as .*',
- ragged.from_row_splits, values, splits)
+ if not context.executing_eagerly():
+ with ops.Graph().as_default():
+ values = constant_op.constant([1, 2, 3], dtypes.int64)
+ with ops.Graph().as_default():
+ splits = constant_op.constant([0, 2, 3], dtypes.int64)
+ self.assertRaisesRegexp(ValueError,
+ '.* must be from the same graph as .*',
+ RaggedTensor.from_row_splits, values, splits)
#=============================================================================
# Ragged Value & Row-Partitioning Tensor Accessors
#=============================================================================
- @test_util.run_deprecated_v1
def testRaggedTensorAccessors_2d(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
- rt1 = ragged.from_row_splits(values, row_splits)
- rt2 = ragged.from_value_rowids(values, value_rowids)
+ rt1 = RaggedTensor.from_row_splits(values, row_splits)
+ rt2 = RaggedTensor.from_value_rowids(values, value_rowids)
for rt in [rt1, rt2]:
- self.assertEqual(
- self.evaluate(rt).tolist(),
- [[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
- self.assertEqual(
- self.evaluate(rt.values).tolist(),
- [b'a', b'b', b'c', b'd', b'e', b'f', b'g'])
+ self.assertRaggedEqual(
+ rt, [[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
+ self.assertAllEqual(rt.values, [b'a', b'b', b'c', b'd', b'e', b'f', b'g'])
self.assertEqual(rt.values.shape.dims[0].value, 7)
- self.assertEqual(
- self.evaluate(ragged.value_rowids(rt)).tolist(),
- [0, 0, 2, 2, 2, 3, 4])
- self.assertEqual(self.evaluate(ragged.nrows(rt)).tolist(), 5)
- self.assertEqual(
- self.evaluate(rt.row_splits).tolist(), [0, 2, 2, 5, 6, 7])
- self.assertEqual(
- self.evaluate(ragged.row_starts(rt)).tolist(), [0, 2, 2, 5, 6])
- self.assertEqual(
- self.evaluate(ragged.row_limits(rt)).tolist(), [2, 2, 5, 6, 7])
- self.assertEqual(
- self.evaluate(ragged.row_lengths(rt)).tolist(), [2, 0, 3, 1, 1])
- self.assertEqual(
- self.evaluate(rt.inner_values).tolist(),
- [b'a', b'b', b'c', b'd', b'e', b'f', b'g'])
- self.assertEqual(
- [self.evaluate(s).tolist() for s in rt.nested_row_splits],
- [[0, 2, 2, 5, 6, 7]])
+ self.assertAllEqual(rt.value_rowids(), [0, 0, 2, 2, 2, 3, 4])
+ self.assertAllEqual(rt.nrows(), 5)
+ self.assertAllEqual(rt.row_splits, [0, 2, 2, 5, 6, 7])
+ self.assertAllEqual(rt.row_starts(), [0, 2, 2, 5, 6])
+ self.assertAllEqual(rt.row_limits(), [2, 2, 5, 6, 7])
+ self.assertAllEqual(rt.row_lengths(), [2, 0, 3, 1, 1])
+ self.assertAllEqual(rt.flat_values,
+ [b'a', b'b', b'c', b'd', b'e', b'f', b'g'])
+ self.assertLen(rt.nested_row_splits, 1)
+ self.assertAllEqual(rt.nested_row_splits[0], [0, 2, 2, 5, 6, 7])
- @test_util.run_deprecated_v1
def testRaggedTensorAccessors_3d_with_ragged_rank_1(self):
values = [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]]
row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
- rt1 = ragged.from_row_splits(values, row_splits)
- rt2 = ragged.from_value_rowids(values, value_rowids)
+ rt1 = RaggedTensor.from_row_splits(values, row_splits)
+ rt2 = RaggedTensor.from_value_rowids(values, value_rowids)
for rt in [rt1, rt2]:
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[[0, 1], [2, 3]], [], [[4, 5], [6, 7], [8, 9]], [[10, 11]],
[[12, 13]]])
self.assertEqual(
- self.evaluate(rt.values).tolist(),
+ self.eval_to_list(rt.values),
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]])
self.assertEqual(rt.values.shape.dims[0].value, 7)
self.assertEqual(
- self.evaluate(ragged.value_rowids(rt)).tolist(),
- [0, 0, 2, 2, 2, 3, 4])
- self.assertEqual(self.evaluate(ragged.nrows(rt)).tolist(), 5)
+ self.eval_to_list(rt.value_rowids()), [0, 0, 2, 2, 2, 3, 4])
+ self.assertEqual(self.eval_to_list(rt.nrows()), 5)
+ self.assertEqual(self.eval_to_list(rt.row_splits), [0, 2, 2, 5, 6, 7])
+ self.assertEqual(self.eval_to_list(rt.row_starts()), [0, 2, 2, 5, 6])
+ self.assertEqual(self.eval_to_list(rt.row_limits()), [2, 2, 5, 6, 7])
+ self.assertEqual(self.eval_to_list(rt.row_lengths()), [2, 0, 3, 1, 1])
self.assertEqual(
- self.evaluate(rt.row_splits).tolist(), [0, 2, 2, 5, 6, 7])
- self.assertEqual(
- self.evaluate(ragged.row_starts(rt)).tolist(), [0, 2, 2, 5, 6])
- self.assertEqual(
- self.evaluate(ragged.row_limits(rt)).tolist(), [2, 2, 5, 6, 7])
- self.assertEqual(
- self.evaluate(ragged.row_lengths(rt)).tolist(), [2, 0, 3, 1, 1])
- self.assertEqual(
- self.evaluate(rt.inner_values).tolist(),
+ self.eval_to_list(rt.flat_values),
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]])
- self.assertEqual(
- [self.evaluate(s).tolist() for s in rt.nested_row_splits],
- [[0, 2, 2, 5, 6, 7]])
+ self.assertEqual([self.eval_to_list(s) for s in rt.nested_row_splits],
+ [[0, 2, 2, 5, 6, 7]])
- @test_util.run_deprecated_v1
def testRaggedTensorAccessors_3d_with_ragged_rank_2(self):
values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
nested_row_splits = [
@@ -685,73 +660,59 @@
constant_op.constant([0, 0, 1, 3, 3], dtypes.int64),
constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
]
- rt1 = ragged.from_nested_row_splits(values, nested_row_splits)
- rt2 = ragged.from_nested_value_rowids(values, nested_value_rowids)
+ rt1 = RaggedTensor.from_nested_row_splits(values, nested_row_splits)
+ rt2 = RaggedTensor.from_nested_value_rowids(values, nested_value_rowids)
for rt in [rt1, rt2]:
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[[b'a', b'b'], []], [[b'c', b'd', b'e']], [], [[b'f'], [b'g']]])
self.assertEqual(
- self.evaluate(rt.values).tolist(),
+ self.eval_to_list(rt.values),
[[b'a', b'b'], [], [b'c', b'd', b'e'], [b'f'], [b'g']])
self.assertEqual(rt.values.shape.dims[0].value, 5)
+ self.assertEqual(self.eval_to_list(rt.value_rowids()), [0, 0, 1, 3, 3])
+ self.assertEqual(self.eval_to_list(rt.nrows()), 4)
+ self.assertEqual(self.eval_to_list(rt.row_splits), [0, 2, 3, 3, 5])
+ self.assertEqual(self.eval_to_list(rt.row_starts()), [0, 2, 3, 3])
+ self.assertEqual(self.eval_to_list(rt.row_limits()), [2, 3, 3, 5])
+ self.assertEqual(self.eval_to_list(rt.row_lengths()), [2, 1, 0, 2])
self.assertEqual(
- self.evaluate(ragged.value_rowids(rt)).tolist(), [0, 0, 1, 3, 3])
- self.assertEqual(self.evaluate(ragged.nrows(rt)).tolist(), 4)
- self.assertEqual(self.evaluate(rt.row_splits).tolist(), [0, 2, 3, 3, 5])
- self.assertEqual(
- self.evaluate(ragged.row_starts(rt)).tolist(), [0, 2, 3, 3])
- self.assertEqual(
- self.evaluate(ragged.row_limits(rt)).tolist(), [2, 3, 3, 5])
- self.assertEqual(
- self.evaluate(ragged.row_lengths(rt)).tolist(), [2, 1, 0, 2])
- self.assertEqual(
- self.evaluate(rt.inner_values).tolist(),
+ self.eval_to_list(rt.flat_values),
[b'a', b'b', b'c', b'd', b'e', b'f', b'g'])
- self.assertEqual(
- [self.evaluate(s).tolist() for s in rt.nested_row_splits],
- [[0, 2, 3, 3, 5], [0, 2, 2, 5, 6, 7]])
-
- def testNRowsWithTensorInput(self):
- dt = constant_op.constant([[1, 2, 3], [4, 5, 6]])
- nrows = ragged.nrows(dt)
- self.assertEqual(self.evaluate(nrows), 2)
-
- def testRowLengthsWithTensorInput(self):
- dt = constant_op.constant([[1, 2, 3], [4, 5, 6]])
- row_lengths = ragged.row_lengths(dt)
- self.assertEqual(self.evaluate(row_lengths).tolist(), [3, 3])
+ self.assertEqual([self.eval_to_list(s) for s in rt.nested_row_splits],
+ [[0, 2, 3, 3, 5], [0, 2, 2, 5, 6, 7]])
#=============================================================================
# RaggedTensor.shape
#=============================================================================
- @test_util.run_deprecated_v1
def testShape(self):
"""Tests for RaggedTensor.shape."""
- rt1 = ragged.from_row_splits(b'a b c d e f g'.split(), [0, 2, 5, 6, 6, 7])
+ rt1 = RaggedTensor.from_row_splits(b'a b c d e f g'.split(),
+ [0, 2, 5, 6, 6, 7])
self.assertEqual(rt1.shape.as_list(), [5, None])
- rt2 = ragged.from_row_splits(
+ rt2 = RaggedTensor.from_row_splits(
[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14]],
[0, 2, 5, 6, 6, 7])
self.assertEqual(rt2.shape.as_list(), [5, None, 2])
- rt3 = ragged.from_row_splits(
+ rt3 = RaggedTensor.from_row_splits(
[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], [0, 2, 2, 3])
self.assertEqual(rt3.shape.as_list(), [3, None, 2, 2])
- rt4 = ragged.from_row_splits(rt3, [0, 1, 3, 3])
+ rt4 = RaggedTensor.from_row_splits(rt3, [0, 1, 3, 3])
self.assertEqual(rt4.shape.as_list(), [3, None, None, 2, 2])
- rt5 = ragged.from_row_splits(
- array_ops.placeholder(dtype=dtypes.string), [0, 2, 3, 5])
- self.assertEqual(rt5.shape.ndims, None)
+ if not context.executing_eagerly():
+ rt5 = RaggedTensor.from_row_splits(
+ array_ops.placeholder(dtype=dtypes.string), [0, 2, 3, 5])
+ self.assertEqual(rt5.shape.ndims, None)
- rt6 = ragged.from_row_splits([1, 2, 3],
- array_ops.placeholder(dtype=dtypes.int64))
- self.assertEqual(rt6.shape.as_list(), [None, None])
+ rt6 = RaggedTensor.from_row_splits(
+ [1, 2, 3], array_ops.placeholder(dtype=dtypes.int64))
+ self.assertEqual(rt6.shape.as_list(), [None, None])
#=============================================================================
# RaggedTensor.__getitem__
@@ -777,15 +738,9 @@
"""
tensor_slice_spec1 = _make_tensor_slice_spec(slice_spec, True)
tensor_slice_spec2 = _make_tensor_slice_spec(slice_spec, False)
- value1 = self.evaluate(rt.__getitem__(slice_spec))
- value2 = self.evaluate(rt.__getitem__(tensor_slice_spec1))
- value3 = self.evaluate(rt.__getitem__(tensor_slice_spec2))
- if hasattr(value1, 'tolist'):
- value1 = value1.tolist()
- if hasattr(value2, 'tolist'):
- value2 = value2.tolist()
- if hasattr(value3, 'tolist'):
- value3 = value3.tolist()
+ value1 = self.eval_to_list(rt.__getitem__(slice_spec))
+ value2 = self.eval_to_list(rt.__getitem__(tensor_slice_spec1))
+ value3 = self.eval_to_list(rt.__getitem__(tensor_slice_spec2))
self.assertEqual(value1, expected, 'slice_spec=%s' % (slice_spec,))
self.assertEqual(value2, expected, 'slice_spec=%s' % (slice_spec,))
self.assertEqual(value3, expected, 'slice_spec=%s' % (slice_spec,))
@@ -861,23 +816,26 @@
(SLICE_BUILDER[:, -2:], [row[-2:] for row in EXAMPLE_RAGGED_TENSOR_2D]),
# TODO(edloper): Add tests for strided slices, once support is added.
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemWithRaggedRank1(self, slice_spec, expected):
"""Test that rt.__getitem__(slice_spec) == expected."""
# Ragged tensor
- rt = ragged.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
- EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
+ rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
+ EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
- self.assertEqual(self.evaluate(rt).tolist(), EXAMPLE_RAGGED_TENSOR_2D)
+ self.assertEqual(self.eval_to_list(rt), EXAMPLE_RAGGED_TENSOR_2D)
self._TestGetItem(rt, slice_spec, expected)
# pylint: disable=invalid-slice-index
@parameterized.parameters(
# Tests for out-of-bound errors
- (SLICE_BUILDER[5], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[-6], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[0, 2], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[3, 0], ValueError, '.*out of bounds.*'),
+ (SLICE_BUILDER[5],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[-6],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[0, 2],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[3, 0],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
# Indexing into an inner ragged dimension
(SLICE_BUILDER[:, 3], ValueError,
@@ -889,8 +847,8 @@
# Tests for type errors
(SLICE_BUILDER[0.5], TypeError, re.escape(array_ops._SLICE_TYPE_ERROR)),
- (SLICE_BUILDER[1:3:0.5], TypeError,
- re.escape(array_ops._SLICE_TYPE_ERROR)),
+ (SLICE_BUILDER[1:3:0.5], TypeError, re.escape(
+ array_ops._SLICE_TYPE_ERROR)),
(SLICE_BUILDER[:, 1:3:0.5], TypeError,
'slice strides must be integers or None'),
(SLICE_BUILDER[:, 0.5:1.5], TypeError,
@@ -903,17 +861,14 @@
(SLICE_BUILDER[..., 0, 0, 0], IndexError,
'Too many indices for RaggedTensor'),
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemErrorsWithRaggedRank1(self, slice_spec, expected,
message):
"""Test that rt.__getitem__(slice_spec) == expected."""
# Ragged tensor
- rt = ragged.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
- EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
- # if sys.version_info[0] == 3:
- # message = 'must be str, not int'
+ rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
+ EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
- self.assertEqual(self.evaluate(rt).tolist(), EXAMPLE_RAGGED_TENSOR_2D)
+ self.assertEqual(self.eval_to_list(rt), EXAMPLE_RAGGED_TENSOR_2D)
self._TestGetItemException(rt, slice_spec, expected, message)
@parameterized.parameters(
@@ -982,13 +937,12 @@
# TODO(edloper): Add tests slicing inner ragged dimensions, one support
# is added.
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemWithRaggedRank2(self, slice_spec, expected):
"""Test that rt.__getitem__(slice_spec) == expected."""
- rt = ragged.from_nested_row_splits(
+ rt = RaggedTensor.from_nested_row_splits(
EXAMPLE_RAGGED_TENSOR_4D_VALUES,
[EXAMPLE_RAGGED_TENSOR_4D_SPLITS1, EXAMPLE_RAGGED_TENSOR_4D_SPLITS2])
- self.assertEqual(self.evaluate(rt).tolist(), EXAMPLE_RAGGED_TENSOR_4D)
+ self.assertEqual(self.eval_to_list(rt), EXAMPLE_RAGGED_TENSOR_4D)
self._TestGetItem(rt, slice_spec, expected)
@parameterized.parameters(
@@ -999,19 +953,22 @@
'Cannot index into an inner ragged dimension.'),
# Test for out-of-bounds errors.
- (SLICE_BUILDER[1, 0], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[0, 0, 3], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[5], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[0, 5], ValueError, '.*out of bounds.*'),
+ (SLICE_BUILDER[1, 0],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[0, 0, 3],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[5],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[0, 5],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemErrorsWithRaggedRank2(self, slice_spec, expected,
message):
"""Test that rt.__getitem__(slice_spec) == expected."""
- rt = ragged.from_nested_row_splits(
+ rt = RaggedTensor.from_nested_row_splits(
EXAMPLE_RAGGED_TENSOR_4D_VALUES,
[EXAMPLE_RAGGED_TENSOR_4D_SPLITS1, EXAMPLE_RAGGED_TENSOR_4D_SPLITS2])
- self.assertEqual(self.evaluate(rt).tolist(), EXAMPLE_RAGGED_TENSOR_4D)
+ self.assertEqual(self.eval_to_list(rt), EXAMPLE_RAGGED_TENSOR_4D)
self._TestGetItemException(rt, slice_spec, expected, message)
@parameterized.parameters(
@@ -1019,21 +976,21 @@
(SLICE_BUILDER[2:], []),
(SLICE_BUILDER[:-3], []),
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemWithEmptyTensor(self, slice_spec, expected):
"""Test that rt.__getitem__(slice_spec) == expected."""
- rt = ragged.from_row_splits([], [0])
+ rt = RaggedTensor.from_row_splits([], [0])
self._TestGetItem(rt, slice_spec, expected)
@parameterized.parameters(
- (SLICE_BUILDER[0], ValueError, '.*out of bounds.*'),
- (SLICE_BUILDER[-1], ValueError, '.*out of bounds.*'),
+ (SLICE_BUILDER[0],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
+ (SLICE_BUILDER[-1],
+ (ValueError, errors.InvalidArgumentError), '.*out of bounds.*'),
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemErrorsWithEmptyTensor(self, slice_spec, expected,
message):
"""Test that rt.__getitem__(slice_spec) == expected."""
- rt = ragged.from_row_splits([], [0])
+ rt = RaggedTensor.from_row_splits([], [0])
self._TestGetItemException(rt, slice_spec, expected, message)
@parameterized.parameters(
@@ -1045,7 +1002,6 @@
(SLICE_BUILDER[0, 1], EXAMPLE_RAGGED_TENSOR_2D[0][1]),
(SLICE_BUILDER[-3, 0], EXAMPLE_RAGGED_TENSOR_2D[-3][0]),
)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemWithPlaceholderShapes(self, slice_spec, expected):
"""Test that rt.__getitem__(slice_spec) == expected."""
# Intentionally use an unknown shape for `splits`, to force the code path
@@ -1053,29 +1009,28 @@
splits = constant_op.constant(
EXAMPLE_RAGGED_TENSOR_2D_SPLITS, dtype=dtypes.int64)
splits = array_ops.placeholder_with_default(splits, None)
- rt = ragged.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES, splits)
- self.assertEqual(self.evaluate(rt).tolist(), EXAMPLE_RAGGED_TENSOR_2D)
+ rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES, splits)
+ self.assertEqual(self.eval_to_list(rt), EXAMPLE_RAGGED_TENSOR_2D)
self._TestGetItem(rt, slice_spec, expected)
@parameterized.parameters(
(SLICE_BUILDER[..., 2], ValueError,
'Ellipsis not supported for unknown shape RaggedTensors'),)
- @test_util.run_deprecated_v1
def testRaggedTensorGetItemErrorsWithPlaceholderShapes(
self, slice_spec, expected, message):
"""Test that rt.__getitem__(slice_spec) == expected."""
- # Intentionally use an unknown shape for `values`.
- values = array_ops.placeholder_with_default([0], None)
- rt = ragged.from_row_splits(values, [0, 1])
- self._TestGetItemException(rt, slice_spec, expected, message)
+ if not context.executing_eagerly():
+ # Intentionally use an unknown shape for `values`.
+ values = array_ops.placeholder_with_default([0], None)
+ rt = RaggedTensor.from_row_splits(values, [0, 1])
+ self._TestGetItemException(rt, slice_spec, expected, message)
- @test_util.run_deprecated_v1
def testGetItemNewAxis(self):
# rt: [[[['a', 'b'], ['c', 'd']], [], [['e', 'f']]], []]
splits1 = [0, 3, 3]
splits2 = [0, 2, 2, 3]
values = constant_op.constant([['a', 'b'], ['c', 'd'], ['e', 'f']])
- rt = ragged.from_nested_row_splits(values, [splits1, splits2])
+ rt = RaggedTensor.from_nested_row_splits(values, [splits1, splits2])
rt_newaxis0 = rt[array_ops.newaxis]
rt_newaxis1 = rt[:, array_ops.newaxis]
rt_newaxis2 = rt[:, :, array_ops.newaxis]
@@ -1083,22 +1038,22 @@
rt_newaxis4 = rt[:, :, :, :, array_ops.newaxis]
self.assertEqual(
- self.evaluate(rt).tolist(),
+ self.eval_to_list(rt),
[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]], []])
self.assertEqual(
- self.evaluate(rt_newaxis0).tolist(),
+ self.eval_to_list(rt_newaxis0),
[[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]], []]])
self.assertEqual(
- self.evaluate(rt_newaxis1).tolist(),
+ self.eval_to_list(rt_newaxis1),
[[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]]], [[]]])
self.assertEqual(
- self.evaluate(rt_newaxis2).tolist(),
+ self.eval_to_list(rt_newaxis2),
[[[[[b'a', b'b'], [b'c', b'd']]], [[]], [[[b'e', b'f']]]], []])
self.assertEqual(
- self.evaluate(rt_newaxis3).tolist(),
+ self.eval_to_list(rt_newaxis3),
[[[[[b'a', b'b']], [[b'c', b'd']]], [], [[[b'e', b'f']]]], []])
self.assertEqual(
- self.evaluate(rt_newaxis4).tolist(),
+ self.eval_to_list(rt_newaxis4),
[[[[[b'a'], [b'b']], [[b'c'], [b'd']]], [], [[[b'e'], [b'f']]]], []])
self.assertEqual(rt.ragged_rank, 2)
@@ -1117,126 +1072,140 @@
#=============================================================================
# RaggedTensor.__str__
#=============================================================================
- @test_util.run_deprecated_v1
def testRaggedTensorStr(self):
- rt1 = ragged.from_row_splits(b'a b c d e f g'.split(), [0, 2, 5, 6, 6, 7])
- expected1 = ('RaggedTensor(values=Tensor("RaggedFromRowSplits/values:0", '
- 'shape=(7,), dtype=string), row_splits='
- 'Tensor("RaggedFromRowSplits/row_splits:0", '
- 'shape=(6,), dtype=int64))')
- self.assertEqual(str(rt1), expected1)
- self.assertEqual(repr(rt1), expected1)
+ values = [b'a', b'b', b'c', b'd', b'e', b'f', b'g']
+ row_splits = [0, 2, 5, 6, 6, 7]
+ rt = RaggedTensor.from_row_splits(values, row_splits)
+ if context.executing_eagerly():
+ expected_str = '<tf.RaggedTensor {}>'.format([[b'a', b'b'],
+ [b'c', b'd', b'e'], [b'f'],
+ [], [b'g']])
+ expected_repr = (
+ 'tf.RaggedTensor(values=tf.Tensor([{}], shape=(7,), dtype=string), '
+ 'row_splits=tf.Tensor([{}], shape=(6,), dtype=int64))'.format(
+ ' '.join(repr(x) for x in values), ' '.join(
+ repr(x) for x in row_splits)))
+ self.assertEqual(str(rt), expected_str)
+ self.assertEqual(repr(rt), expected_repr)
+ else:
+ expected_repr = (
+ 'tf.RaggedTensor(values=Tensor("RaggedFromRowSplits/values:0", '
+ 'shape=(7,), dtype=string), row_splits='
+ 'Tensor("RaggedFromRowSplits/row_splits:0", '
+ 'shape=(6,), dtype=int64))')
+ self.assertEqual(repr(rt), expected_repr)
+ self.assertEqual(str(rt), expected_repr)
def testRaggedTensorValueStr(self):
- rt = ragged.RaggedTensorValue(
- values=np.array(b'a b c d e f g'.split()),
- row_splits=np.array([0, 2, 5, 6, 6, 7], dtype=np.int64))
- if sys.version_info[0] == 2:
- self.assertEqual(' '.join(str(rt).split()),
- (r"<RaggedTensorValue [['a', 'b'], ['c', 'd', 'e'], "
- "['f'], [], ['g']]>"))
- self.assertEqual(
- ' '.join(repr(rt).split()),
- (r"RaggedTensorValue(values=array(['a', 'b', 'c', 'd', "
- "'e', 'f', 'g'], dtype='|S1'), row_splits=array([0, 2, 5,"
- ' 6, 6, 7]))'))
- else:
- self.assertEqual(
- ' '.join(str(rt).split()),
- (r"<RaggedTensorValue [[b'a', b'b'], [b'c', b'd', b'e'], "
- "[b'f'], [], [b'g']]>"))
- self.assertEqual(
- ' '.join(repr(rt).split()),
- (r"RaggedTensorValue(values=array([b'a', b'b', b'c', b'd', "
- "b'e', b'f', b'g'], dtype='|S1'), row_splits=array([0, 2, 5,"
- ' 6, 6, 7]))'))
+ values = [b'a', b'b', b'c', b'd', b'e', b'f', b'g']
+ row_splits = [0, 2, 5, 6, 6, 7]
+ rt = ragged_tensor_value.RaggedTensorValue(
+ np.array(values), np.array(row_splits, dtype=np.int64))
+ expected_str = '<tf.RaggedTensorValue {}>'.format([[b'a', b'b'],
+ [b'c', b'd', b'e'],
+ [b'f'], [], [b'g']])
+ expected_repr = ("tf.RaggedTensorValue(values=array({}, dtype='|S1'), "
+ 'row_splits=array({}))'.format(values, row_splits))
+ self.assertEqual(' '.join(str(rt).split()), expected_str)
+ self.assertEqual(' '.join(repr(rt).split()), expected_repr)
#=============================================================================
- # RaggedTensor.with_values() and RaggedTensor.with_inner_values().
+ # RaggedTensor.with_values() and RaggedTensor.with_flat_values().
#=============================================================================
- @test_util.run_deprecated_v1
def testWithValues(self):
- rt1 = ragged.constant([[1, 2], [3, 4, 5], [6], [], [7]])
- rt2 = ragged.constant([[[1, 2], [3, 4, 5]], [[6]], [], [[], [7]]])
+ rt1 = ragged_factory_ops.constant([[1, 2], [3, 4, 5], [6], [], [7]])
+ rt2 = ragged_factory_ops.constant([[[1, 2], [3, 4, 5]], [[6]], [], [[],
+ [7]]])
rt1_plus_10 = rt1.with_values(rt1.values + 10)
- rt2_times_10 = rt2.with_inner_values(rt2.inner_values * 10)
+ rt2_times_10 = rt2.with_flat_values(rt2.flat_values * 10)
rt1_expanded = rt1.with_values(array_ops.expand_dims(rt1.values, axis=1))
self.assertEqual(
- self.evaluate(rt1_plus_10).tolist(),
+ self.eval_to_list(rt1_plus_10),
[[11, 12], [13, 14, 15], [16], [], [17]])
self.assertEqual(
- self.evaluate(rt2_times_10).tolist(),
+ self.eval_to_list(rt2_times_10),
[[[10, 20], [30, 40, 50]], [[60]], [], [[], [70]]])
self.assertEqual(
- self.evaluate(rt1_expanded).tolist(),
+ self.eval_to_list(rt1_expanded),
[[[1], [2]], [[3], [4], [5]], [[6]], [], [[7]]])
#=============================================================================
# Session.run
#=============================================================================
- @test_util.run_deprecated_v1
def testSessionRun(self):
- rt1 = ragged.constant([[1, 2, 3], [4]])
- rt2 = ragged.constant([[[], [1, 2]], [[3]]])
+ if context.executing_eagerly():
+ return
+
+ rt1 = ragged_factory_ops.constant([[1, 2, 3], [4]])
+ rt2 = ragged_factory_ops.constant([[[], [1, 2]], [[3]]])
with self.test_session() as session:
result = session.run({'rt1': rt1, 'rt2': rt2})
self.assertCountEqual(sorted(result.keys()), ['rt1', 'rt2'])
- self.assertEqual(result['rt1'].tolist(), [[1, 2, 3], [4]])
- self.assertEqual(result['rt2'].tolist(), [[[], [1, 2]], [[3]]])
+ self.assertEqual(result['rt1'].to_list(), [[1, 2, 3], [4]])
+ self.assertEqual(result['rt2'].to_list(), [[[], [1, 2]], [[3]]])
- @test_util.run_deprecated_v1
def testSessionRunFeed(self):
- rt1 = ragged.from_row_splits(
+ if context.executing_eagerly():
+ return
+
+ rt1 = RaggedTensor.from_row_splits(
array_ops.placeholder(dtypes.int32),
array_ops.placeholder(dtypes.int64))
- rt2 = ragged.from_nested_row_splits(
- array_ops.placeholder(dtypes.int32),
- [array_ops.placeholder(dtypes.int64),
- array_ops.placeholder(dtypes.int64)])
+ rt2 = RaggedTensor.from_nested_row_splits(
+ array_ops.placeholder(dtypes.int32), [
+ array_ops.placeholder(dtypes.int64),
+ array_ops.placeholder(dtypes.int64)
+ ])
- rt1_feed_val = ragged.constant_value([[1, 2, 3], [4]])
- rt2_feed_val = ragged.constant_value([[[], [1, 2]], [[3]]])
+ rt1_feed_val = ragged_factory_ops.constant_value([[1, 2, 3], [4]])
+ rt2_feed_val = ragged_factory_ops.constant_value([[[], [1, 2]], [[3]]])
with self.test_session() as session:
- result = session.run({'rt1': rt1, 'rt2': rt2},
- feed_dict={rt1: rt1_feed_val,
- rt2: rt2_feed_val})
+ result = session.run({
+ 'rt1': rt1,
+ 'rt2': rt2
+ },
+ feed_dict={
+ rt1: rt1_feed_val,
+ rt2: rt2_feed_val
+ })
self.assertCountEqual(sorted(result.keys()), ['rt1', 'rt2'])
- self.assertEqual(result['rt1'].tolist(), [[1, 2, 3], [4]])
- self.assertEqual(result['rt2'].tolist(), [[[], [1, 2]], [[3]]])
+ self.assertEqual(result['rt1'].to_list(), [[1, 2, 3], [4]])
+ self.assertEqual(result['rt2'].to_list(), [[[], [1, 2]], [[3]]])
- @test_util.run_deprecated_v1
def testSessionPartialRunFeed(self):
+ if context.executing_eagerly():
+ return
+
# Placeholder inputs.
- a = ragged.from_row_splits(
+ a = RaggedTensor.from_row_splits(
array_ops.placeholder(dtypes.int32, shape=[None], name='a.values'),
array_ops.placeholder(dtypes.int64, name='a.row_splits'))
- b = ragged.from_row_splits(
+ b = RaggedTensor.from_row_splits(
array_ops.placeholder(dtypes.int32, shape=[None], name='b.values'),
array_ops.placeholder(dtypes.int64, name='b.row_splits'))
c = array_ops.placeholder(dtypes.int32, shape=[], name='c')
# Feed values for placeholder inputs.
- a_val = ragged.constant_value([[1, 2, 3], [4]])
- b_val = ragged.constant_value([[5, 4, 3], [2]])
+ a_val = ragged_factory_ops.constant_value([[1, 2, 3], [4]])
+ b_val = ragged_factory_ops.constant_value([[5, 4, 3], [2]])
c_val = 3
# Compute some values.
- r1 = ragged.reduce_sum(a * b, axis=1)
- r2 = ragged.reduce_sum(a + c, axis=1)
+ r1 = ragged_math_ops.reduce_sum(a * b, axis=1)
+ r2 = ragged_math_ops.reduce_sum(a + c, axis=1)
with self.test_session() as session:
handle = session.partial_run_setup([r1, r2], [a, b, c])
res1 = session.partial_run(handle, r1, feed_dict={a: a_val, b: b_val})
- self.assertEqual(res1.tolist(), [22, 8])
+ self.assertAllEqual(res1, [22, 8])
res2 = session.partial_run(handle, r2, feed_dict={c: c_val})
- self.assertEqual(res2.tolist(), [15, 7])
-
+ self.assertAllEqual(res2, [15, 7])
if __name__ == '__main__':
googletest.main()
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_value.py b/tensorflow/python/ops/ragged/ragged_tensor_value.py
index 39d3249..c5e498e 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_value.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_value.py
@@ -20,11 +20,17 @@
import numpy as np
+from tensorflow.python.util.tf_export import tf_export
+
+@tf_export(v1=["ragged.RaggedTensorValue"])
class RaggedTensorValue(object):
"""Represents the value of a `RaggedTensor`.
- See `RaggedTensor` for a description of ragged tensors.
+ Warning: `RaggedTensorValue` should only be used in graph mode; in
+ eager mode, the `tf.RaggedTensor` class contains its value directly.
+
+ See `tf.RaggedTensor` for a description of ragged tensors.
"""
def __init__(self, values, row_splits):
@@ -53,7 +59,7 @@
doc="""The numpy dtype of values in this tensor.""")
@property
- def inner_values(self):
+ def flat_values(self):
"""The innermost `values` array for this ragged tensor value."""
rt_values = self.values
while isinstance(rt_values, RaggedTensorValue):
@@ -82,15 +88,18 @@
return (self._row_splits.shape[0] - 1,) + (None,) + self._values.shape[1:]
def __str__(self):
- return "<RaggedTensorValue %s>" % self.tolist()
+ return "<tf.RaggedTensorValue %s>" % self.to_list()
def __repr__(self):
- return "RaggedTensorValue(values=%r, row_splits=%r)" % (self._values,
- self._row_splits)
+ return "tf.RaggedTensorValue(values=%r, row_splits=%r)" % (self._values,
+ self._row_splits)
- def tolist(self):
+ def to_list(self):
"""Returns this ragged tensor value as a nested Python list."""
- values_as_list = self._values.tolist()
+ if isinstance(self._values, RaggedTensorValue):
+ values_as_list = self._values.to_list()
+ else:
+ values_as_list = self._values.tolist()
return [
values_as_list[self._row_splits[i]:self._row_splits[i + 1]]
for i in range(len(self._row_splits) - 1)
diff --git a/tensorflow/python/ops/ragged/ragged_test_util.py b/tensorflow/python/ops/ragged/ragged_test_util.py
new file mode 100644
index 0000000..dcbab30
--- /dev/null
+++ b/tensorflow/python/ops/ragged/ragged_test_util.py
@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+
+# pylint: disable=invalid-name
+"""Test utils for tensorflow RaggedTensors."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_value
+
+
+class RaggedTensorTestCase(test_util.TensorFlowTestCase):
+ """Base class for RaggedTensor test cases."""
+
+ def _GetPyList(self, a):
+ """Converts a to a nested python list."""
+ if isinstance(a, ragged_tensor.RaggedTensor):
+ return self.evaluate(a).to_list()
+ elif isinstance(a, ops.Tensor):
+ a = self.evaluate(a)
+ return a.tolist() if isinstance(a, np.ndarray) else a
+ elif isinstance(a, np.ndarray):
+ return a.tolist()
+ elif isinstance(a, ragged_tensor_value.RaggedTensorValue):
+ return a.to_list()
+ else:
+ return np.array(a).tolist()
+
+ def assertRaggedEqual(self, a, b):
+ """Asserts that two potentially ragged tensors are equal."""
+ a_list = self._GetPyList(a)
+ b_list = self._GetPyList(b)
+ self.assertEqual(a_list, b_list)
+
+ if not (isinstance(a, (list, tuple)) or isinstance(b, (list, tuple))):
+ a_ragged_rank = a.ragged_rank if ragged_tensor.is_ragged(a) else 0
+ b_ragged_rank = b.ragged_rank if ragged_tensor.is_ragged(b) else 0
+ self.assertEqual(a_ragged_rank, b_ragged_rank)
+
+ def assertRaggedAlmostEqual(self, a, b, places=7):
+ a_list = self._GetPyList(a)
+ b_list = self._GetPyList(b)
+ self.assertNestedListAlmostEqual(a_list, b_list, places, context='value')
+
+ if not (isinstance(a, (list, tuple)) or isinstance(b, (list, tuple))):
+ a_ragged_rank = a.ragged_rank if ragged_tensor.is_ragged(a) else 0
+ b_ragged_rank = b.ragged_rank if ragged_tensor.is_ragged(b) else 0
+ self.assertEqual(a_ragged_rank, b_ragged_rank)
+
+ def assertNestedListAlmostEqual(self, a, b, places=7, context='value'):
+ self.assertEqual(type(a), type(b))
+ if isinstance(a, (list, tuple)):
+ self.assertLen(a, len(b), 'Length differs for %s' % context)
+ for i in range(len(a)):
+ self.assertNestedListAlmostEqual(a[i], b[i], places,
+ '%s[%s]' % (context, i))
+ else:
+ self.assertAlmostEqual(
+ a, b, places,
+ '%s != %s within %s places at %s' % (a, b, places, context))
+
+ def eval_to_list(self, tensor):
+ value = self.evaluate(tensor)
+ if ragged_tensor.is_ragged(value):
+ return value.to_list()
+ elif isinstance(value, np.ndarray):
+ return value.tolist()
+ else:
+ return value
+
+ def _eval_tensor(self, tensor):
+ if ragged_tensor.is_ragged(tensor):
+ return ragged_tensor_value.RaggedTensorValue(
+ self._eval_tensor(tensor.values),
+ self._eval_tensor(tensor.row_splits))
+ else:
+ return test_util.TensorFlowTestCase._eval_tensor(self, tensor)
diff --git a/tensorflow/python/ops/ragged/ragged_tile_op_test.py b/tensorflow/python/ops/ragged/ragged_tile_op_test.py
index f335b15..8c03b16 100644
--- a/tensorflow/python/ops/ragged/ragged_tile_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tile_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.tile."""
+"""Tests for ragged_array_ops.tile."""
from __future__ import absolute_import
from __future__ import division
@@ -26,10 +26,13 @@
from tensorflow.python.ops import array_ops
from tensorflow.python.ops.ragged import ragged_array_ops
from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedTileOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTileOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters([
#=========================================================================
@@ -181,7 +184,6 @@
[[[5], [6]]]]),
]) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedTile(self,
descr,
rt_input,
@@ -207,10 +209,8 @@
self.assertEqual(tiled.shape.ndims, rt.shape.ndims)
if multiples_tensor is const_multiples:
self.assertEqual(tiled.shape.as_list(), expected_shape)
- with self.test_session():
- self.assertEqual(tiled.eval().tolist(), expected)
+ self.assertRaggedEqual(tiled, expected)
- @test_util.run_deprecated_v1
def testRaggedTileWithTensorInput(self):
# When the input is a `Tensor`, ragged_tile just delegates to tf.tile.
dt = constant_op.constant([[1, 2], [3, 4]])
@@ -218,8 +218,7 @@
expected = [[1, 2, 1, 2], [3, 4, 3, 4],
[1, 2, 1, 2], [3, 4, 3, 4],
[1, 2, 1, 2], [3, 4, 3, 4]] # pyformat: disable
- with self.test_session():
- self.assertEqual(tiled.eval().tolist(), expected)
+ self.assertRaggedEqual(tiled, expected)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
index 69b31ad..92959a9 100644
--- a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
@@ -18,183 +18,182 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToSparseOpTest(ragged_test_util.RaggedTensorTestCase):
- @test_util.run_deprecated_v1
def testDocStringExample(self):
- rt = ragged.constant([[1, 2, 3], [4], [], [5, 6]])
- st = ragged.to_sparse(rt)
- expected = ('SparseTensorValue(indices='
- 'array([[0, 0], [0, 1], [0, 2], [1, 0], [3, 0], [3, 1]]), '
- 'values=array([1, 2, 3, 4, 5, 6], dtype=int32), '
- 'dense_shape=array([4, 3]))')
- with self.test_session():
- self.assertEqual(' '.join(repr(st.eval()).split()), expected)
+ rt = ragged_factory_ops.constant([[1, 2, 3], [4], [], [5, 6]])
+ st = self.evaluate(rt.to_sparse())
+ self.assertAllEqual(st.indices,
+ [[0, 0], [0, 1], [0, 2], [1, 0], [3, 0], [3, 1]])
+ self.assertAllEqual(st.values, [1, 2, 3, 4, 5, 6])
+ self.assertAllEqual(st.dense_shape, [4, 3])
- @test_util.run_deprecated_v1
def test2DRaggedTensorWithOneRaggedDimension(self):
- rt = ragged.constant([['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']])
- with self.test_session():
- st = ragged.to_sparse(rt).eval()
- self.assertAllEqual(
- st.indices, [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 0], [4, 0]])
- self.assertAllEqual(st.values, b'a b c d e f g'.split())
- self.assertAllEqual(st.dense_shape, [5, 3])
+ rt = ragged_factory_ops.constant([['a', 'b'], ['c', 'd', 'e'], ['f'], [],
+ ['g']])
+ st = self.evaluate(rt.to_sparse())
+ self.assertAllEqual(
+ st.indices, [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 0], [4, 0]])
+ self.assertAllEqual(st.values, b'a b c d e f g'.split())
+ self.assertAllEqual(st.dense_shape, [5, 3])
- @test_util.run_deprecated_v1
def test3DRaggedTensorWithOneRaggedDimension(self):
- rt = ragged.constant([[[1, 2], [3, 4]], [[5, 6], [7, 8], [9, 10]],
- [[11, 12]], [], [[13, 14]]],
- ragged_rank=1)
- with self.test_session():
- st = ragged.to_sparse(rt).eval()
- self.assertAllEqual(
- st.indices, [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0],
- [1, 0, 1], [1, 1, 0], [1, 1, 1], [1, 2, 0], [1, 2, 1],
- [2, 0, 0], [2, 0, 1], [4, 0, 0], [4, 0, 1]])
- self.assertAllEqual(st.values,
- [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
- self.assertAllEqual(st.dense_shape, [5, 3, 2])
+ rt = ragged_factory_ops.constant(
+ [[[1, 2], [3, 4]], [[5, 6], [7, 8], [9, 10]], [[11, 12]], [], [[13, 14]]
+ ],
+ ragged_rank=1)
+ st = self.evaluate(rt.to_sparse())
+ self.assertAllEqual(st.indices,
+ [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0],
+ [1, 0, 1], [1, 1, 0], [1, 1, 1], [1, 2, 0], [1, 2, 1],
+ [2, 0, 0], [2, 0, 1], [4, 0, 0], [4, 0, 1]])
+ self.assertAllEqual(st.values,
+ [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
+ self.assertAllEqual(st.dense_shape, [5, 3, 2])
- @test_util.run_deprecated_v1
def test4DRaggedTensorWithOneRaggedDimension(self):
- rt = ragged.constant(
+ rt = ragged_factory_ops.constant(
[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], [], [[[9, 10], [11, 12]]]],
ragged_rank=1)
- with self.test_session():
- st = ragged.to_sparse(rt).eval()
- self.assertAllEqual(st.values, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
- self.assertAllEqual(
- st.indices,
- [
- [0, 0, 0, 0], # index for value=1
- [0, 0, 0, 1], # index for value=2
- [0, 0, 1, 0], # index for value=3
- [0, 0, 1, 1], # index for value=4
- [0, 1, 0, 0], # index for value=5
- [0, 1, 0, 1], # index for value=6
- [0, 1, 1, 0], # index for value=7
- [0, 1, 1, 1], # index for value=8
- [2, 0, 0, 0], # index for value=9
- [2, 0, 0, 1], # index for value=10
- [2, 0, 1, 0], # index for value=11
- [2, 0, 1, 1], # index for value=12
- ])
- self.assertAllEqual(st.dense_shape, [3, 2, 2, 2])
+ st = self.evaluate(rt.to_sparse())
+ self.assertAllEqual(st.values, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+ self.assertAllEqual(
+ st.indices,
+ [
+ [0, 0, 0, 0], # index for value=1
+ [0, 0, 0, 1], # index for value=2
+ [0, 0, 1, 0], # index for value=3
+ [0, 0, 1, 1], # index for value=4
+ [0, 1, 0, 0], # index for value=5
+ [0, 1, 0, 1], # index for value=6
+ [0, 1, 1, 0], # index for value=7
+ [0, 1, 1, 1], # index for value=8
+ [2, 0, 0, 0], # index for value=9
+ [2, 0, 0, 1], # index for value=10
+ [2, 0, 1, 0], # index for value=11
+ [2, 0, 1, 1], # index for value=12
+ ])
+ self.assertAllEqual(st.dense_shape, [3, 2, 2, 2])
- @test_util.run_deprecated_v1
def test4DRaggedTensorWithTwoRaggedDimensions(self):
- rt = ragged.constant([[[[1, 2], [3, 4]], [[5, 6], [7, 8], [9, 10]]],
- [[[11, 12]], [], [[13, 14]]], []],
- ragged_rank=2)
- with self.test_session():
- st = ragged.to_sparse(rt).eval()
- self.assertAllEqual(
- st.indices,
- [
- [0, 0, 0, 0], # index for value=1
- [0, 0, 0, 1], # index for value=2
- [0, 0, 1, 0], # index for value=3
- [0, 0, 1, 1], # index for value=4
- [0, 1, 0, 0], # index for value=5
- [0, 1, 0, 1], # index for value=6
- [0, 1, 1, 0], # index for value=7
- [0, 1, 1, 1], # index for value=8
- [0, 1, 2, 0], # index for value=9
- [0, 1, 2, 1], # index for value=10
- [1, 0, 0, 0], # index for value=11
- [1, 0, 0, 1], # index for value=12
- [1, 2, 0, 0], # index for value=13
- [1, 2, 0, 1], # index for value=14
- ])
- self.assertAllEqual(st.values,
- [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
- self.assertAllEqual(st.dense_shape, [3, 3, 3, 2])
+ rt = ragged_factory_ops.constant(
+ [[[[1, 2], [3, 4]], [[5, 6], [7, 8], [9, 10]]],
+ [[[11, 12]], [], [[13, 14]]], []],
+ ragged_rank=2)
+ st = self.evaluate(rt.to_sparse())
+ self.assertAllEqual(
+ st.indices,
+ [
+ [0, 0, 0, 0], # index for value=1
+ [0, 0, 0, 1], # index for value=2
+ [0, 0, 1, 0], # index for value=3
+ [0, 0, 1, 1], # index for value=4
+ [0, 1, 0, 0], # index for value=5
+ [0, 1, 0, 1], # index for value=6
+ [0, 1, 1, 0], # index for value=7
+ [0, 1, 1, 1], # index for value=8
+ [0, 1, 2, 0], # index for value=9
+ [0, 1, 2, 1], # index for value=10
+ [1, 0, 0, 0], # index for value=11
+ [1, 0, 0, 1], # index for value=12
+ [1, 2, 0, 0], # index for value=13
+ [1, 2, 0, 1], # index for value=14
+ ])
+ self.assertAllEqual(st.values,
+ [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
+ self.assertAllEqual(st.dense_shape, [3, 3, 3, 2])
def testShape(self):
- rt = ragged.constant([[1, 2], [3, 4, 5], [6], [], [7]])
- st = ragged.to_sparse(rt)
+ rt = ragged_factory_ops.constant([[1, 2], [3, 4, 5], [6], [], [7]])
+ st = rt.to_sparse()
self.assertEqual(st.indices.shape.as_list(), [7, 2])
self.assertEqual(st.values.shape.as_list(), [7])
self.assertEqual(st.dense_shape.shape.as_list(), [2])
- rt = ragged.constant([[[1, 2]], [], [[3, 4]], []], ragged_rank=1)
- st = ragged.to_sparse(rt)
+ rt = ragged_factory_ops.constant([[[1, 2]], [], [[3, 4]], []],
+ ragged_rank=1)
+ st = rt.to_sparse()
self.assertEqual(st.indices.shape.as_list(), [4, 3])
self.assertEqual(st.values.shape.as_list(), [4])
self.assertEqual(st.dense_shape.shape.as_list(), [3])
- rt = ragged.constant([[[1], [2, 3, 4, 5, 6, 7]], [[]]])
- st = ragged.to_sparse(rt)
+ rt = ragged_factory_ops.constant([[[1], [2, 3, 4, 5, 6, 7]], [[]]])
+ st = rt.to_sparse()
self.assertEqual(st.indices.shape.as_list(), [7, 3])
self.assertEqual(st.values.shape.as_list(), [7])
self.assertEqual(st.dense_shape.shape.as_list(), [3])
- @test_util.run_deprecated_v1
def testKernelErrors(self):
# An empty vector, defined using a placeholder to ensure that we can't
# determine that it's invalid at graph-construction time.
empty_vector = array_ops.placeholder_with_default(
array_ops.zeros([0], dtypes.int64), shape=None)
- bad_rt1 = ragged.from_row_splits(row_splits=[2, 3], values=[1, 2, 3])
- with self.test_session():
- bad_split0_error = r'First value of ragged splits must be 0.*'
- self.assertRaisesRegexp(errors.InvalidArgumentError, bad_split0_error,
- ragged.to_sparse(bad_rt1).eval)
+ bad_rt1 = ragged_tensor.RaggedTensor.from_row_splits(
+ row_splits=[2, 3], values=[1, 2, 3])
+ bad_split0 = r'First value of ragged splits must be 0.*'
+ with self.assertRaisesRegexp(errors.InvalidArgumentError, bad_split0):
+ self.evaluate(bad_rt1.to_sparse())
- bad_rt2 = ragged.from_row_splits(row_splits=[0, 5], values=empty_vector)
- bad_rt3 = ragged.from_row_splits(
+ bad_rt2 = ragged_tensor.RaggedTensor.from_row_splits(
+ row_splits=[0, 5], values=empty_vector)
+ bad_rt3 = ragged_tensor.RaggedTensor.from_row_splits(
row_splits=[0, 1],
- values=ragged.from_row_splits(row_splits=[0, 5], values=empty_vector))
- with self.test_session():
- split_mismatch1_error = r'Final value of ragged splits must match.*'
- for rt in [bad_rt2, bad_rt3]:
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- split_mismatch1_error,
- ragged.to_sparse(rt).eval)
+ values=ragged_tensor.RaggedTensor.from_row_splits(
+ row_splits=[0, 5], values=empty_vector))
+ split_mismatch1_error = r'Final value of ragged splits must match.*'
+ for rt in [bad_rt2, bad_rt3]:
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ split_mismatch1_error):
+ self.evaluate(rt.to_sparse())
- bad_rt4 = ragged.from_row_splits(
+ bad_rt4 = ragged_tensor.RaggedTensor.from_row_splits(
row_splits=[0, 5],
- values=ragged.from_row_splits(row_splits=[0], values=empty_vector))
- with self.test_session():
- split_mismatch2_error = r'Final value of ragged splits must match.*'
- self.assertRaisesRegexp(errors.InvalidArgumentError,
- split_mismatch2_error,
- ragged.to_sparse(bad_rt4).eval)
+ values=ragged_tensor.RaggedTensor.from_row_splits(
+ row_splits=[0], values=empty_vector))
+ split_mismatch2_error = r'Final value of ragged splits must match.*'
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ split_mismatch2_error):
+ self.evaluate(bad_rt4.to_sparse())
- bad_rt5 = ragged.from_row_splits(row_splits=empty_vector, values=[])
- with self.test_session():
- empty_splits_error = (r'ragged splits may not be empty.*')
- self.assertRaisesRegexp(errors.InvalidArgumentError, empty_splits_error,
- ragged.to_sparse(bad_rt5).eval)
+ bad_rt5 = ragged_tensor.RaggedTensor.from_row_splits(
+ row_splits=empty_vector, values=[])
+ empty_splits_error = (r'ragged splits may not be empty.*')
+ with self.assertRaisesRegexp(errors.InvalidArgumentError,
+ empty_splits_error):
+ self.evaluate(bad_rt5.to_sparse())
- @test_util.run_deprecated_v1
def testGradient(self):
+ if context.executing_eagerly():
+ return
# rt1.shape == rt2.shape == [2, (D2), (D3), 2].
- rt1 = ragged.constant([[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0]]]],
- ragged_rank=2)
- rt2 = ragged.constant([[[[9.0, 8.0], [7.0, 6.0]], [[5.0, 4.0]]]],
- ragged_rank=2)
- rt = ragged.map_inner_values(math_ops.add, rt1, rt2 * 2.0)
- st = ragged.to_sparse(rt)
+ rt1 = ragged_factory_ops.constant(
+ [[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0]]]], ragged_rank=2)
+ rt2 = ragged_factory_ops.constant(
+ [[[[9.0, 8.0], [7.0, 6.0]], [[5.0, 4.0]]]], ragged_rank=2)
+ rt = ragged_functional_ops.map_flat_values(math_ops.add, rt1, rt2 * 2.0)
+ st = rt.to_sparse()
- g1, g2 = gradients_impl.gradients(st.values, [rt1.inner_values,
- rt2.inner_values])
+ g1, g2 = gradients_impl.gradients(st.values,
+ [rt1.flat_values, rt2.flat_values])
print(g1, g2)
- with self.test_session():
- self.assertEqual(g1.eval().tolist(), [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]])
- self.assertEqual(g2.eval().tolist(), [[2.0, 2.0], [2.0, 2.0], [2.0, 2.0]])
+ self.assertRaggedEqual(g1, [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]])
+ self.assertRaggedEqual(g2, [[2.0, 2.0], [2.0, 2.0], [2.0, 2.0]])
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_to_tensor_op_test.py b/tensorflow/python/ops/ragged/ragged_to_tensor_op_test.py
index 77499b9..ac75456 100644
--- a/tensorflow/python/ops/ragged/ragged_to_tensor_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_to_tensor_op_test.py
@@ -23,24 +23,20 @@
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedTensorToTensorOpTest(test_util.TensorFlowTestCase,
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToTensorOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
- @test_util.run_deprecated_v1
def testDocStringExamples(self):
"""Example from ragged_to_tensor.__doc__."""
- rt = ragged.constant([[9, 8, 7], [], [6, 5], [4]])
- dt = ragged.to_tensor(rt)
- with self.test_session():
- self.assertEqual(str(dt.eval()),
- '[[9 8 7]\n'
- ' [0 0 0]\n'
- ' [6 5 0]\n'
- ' [4 0 0]]') # pyformat: disable
+ rt = ragged_factory_ops.constant([[9, 8, 7], [], [6, 5], [4]])
+ dt = rt.to_tensor()
+ self.assertAllEqual(dt, [[9, 8, 7], [0, 0, 0], [6, 5, 0], [4, 0, 0]])
@parameterized.parameters(
{
@@ -98,23 +94,21 @@
'expected': [[[[1], [2]], [[9], [9]], [[3], [9]]]],
},
)
- @test_util.run_deprecated_v1
def testRaggedTensorToTensor(self,
rt_input,
expected,
ragged_rank=None,
default=None,
expected_shape=None):
- rt = ragged.constant(rt_input, ragged_rank=ragged_rank)
- dt = ragged.to_tensor(rt, default)
- self.assertEqual(type(dt), ops.Tensor)
+ rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
+ dt = rt.to_tensor(default)
+ self.assertIsInstance(dt, ops.Tensor)
self.assertEqual(rt.dtype, dt.dtype)
self.assertTrue(dt.shape.is_compatible_with(rt.shape))
- with self.test_session():
- self.assertEqual(dt.eval().tolist(), expected)
- if expected_shape is not None:
- dt_shape = array_ops.shape(dt)
- self.assertEqual(dt_shape.eval().tolist(), expected_shape)
+ self.assertAllEqual(self.eval_to_list(dt), expected)
+ if expected_shape is not None:
+ dt_shape = array_ops.shape(dt)
+ self.assertAllEqual(dt_shape, expected_shape)
@parameterized.parameters(
{
@@ -131,14 +125,13 @@
{
'rt_input': [[1, 2, 3]],
'default': 'a',
- 'error': (TypeError, "Expected int32, got 'a' of type 'str' instead"),
+ 'error': (TypeError, '.*'),
},
)
- @test_util.run_deprecated_v1
def testError(self, rt_input, default, error, ragged_rank=None):
- rt = ragged.constant(rt_input, ragged_rank=ragged_rank)
+ rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
with self.assertRaisesRegexp(error[0], error[1]):
- ragged.to_tensor(rt, default)
+ rt.to_tensor(default)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_util_test.py b/tensorflow/python/ops/ragged/ragged_util_test.py
index c24ea65..ab5436a 100644
--- a/tensorflow/python/ops/ragged/ragged_util_test.py
+++ b/tensorflow/python/ops/ragged/ragged_util_test.py
@@ -24,6 +24,7 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.ops.ragged import ragged_util
from tensorflow.python.platform import googletest
@@ -41,7 +42,9 @@
for i in range(4)]
-class RaggedRepeatTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedUtilTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters([
# Docstring examples
@@ -89,8 +92,7 @@
])
def testRepeat(self, data, repeats, expected, axis=None):
result = ragged_util.repeat(data, repeats, axis)
- with self.test_session():
- self.assertEqual(result.eval().tolist(), expected)
+ self.assertAllEqual(result, expected)
@parameterized.parameters([
dict(mode=mode, **args)
@@ -155,8 +157,7 @@
repeats = array_ops.placeholder_with_default(repeats, None)
result = ragged_util.repeat(data, repeats, axis)
- with self.test_session():
- self.assertEqual(result.eval().tolist(), expected.tolist())
+ self.assertAllEqual(result, expected)
@parameterized.parameters([
dict(
diff --git a/tensorflow/python/ops/ragged/ragged_where_op_test.py b/tensorflow/python/ops/ragged/ragged_where_op_test.py
index de83a54..3dd9565 100644
--- a/tensorflow/python/ops/ragged/ragged_where_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_where_op_test.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Tests for ragged.where."""
+"""Tests for ragged_array_ops.where."""
from __future__ import absolute_import
from __future__ import division
@@ -21,29 +21,39 @@
from absl.testing import parameterized
from tensorflow.python.framework import test_util
-from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_test_util
from tensorflow.python.platform import googletest
-class RaggedWhereOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedWhereOpTest(ragged_test_util.RaggedTensorTestCase,
+ parameterized.TestCase):
@parameterized.parameters([
#=========================================================================
# Docstring Examples
#=========================================================================
dict( # shape=[D1, (D2)]
- condition=ragged.constant_value([[True, False, True], [False, True]]),
+ condition=ragged_factory_ops.constant_value(
+ [[True, False, True], [False, True]]),
expected=[[0, 0], [0, 2], [1, 1]]),
dict( # shape=[D1, (D2)]
- condition=ragged.constant_value([[True, False, True], [False, True]]),
- x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E']]),
- y=ragged.constant_value([['a', 'b', 'c'], ['d', 'e']]),
- expected=ragged.constant_value([[b'A', b'b', b'C'], [b'd', b'E']])),
+ condition=ragged_factory_ops.constant_value(
+ [[True, False, True], [False, True]]),
+ x=ragged_factory_ops.constant_value(
+ [['A', 'B', 'C'], ['D', 'E']]),
+ y=ragged_factory_ops.constant_value(
+ [['a', 'b', 'c'], ['d', 'e']]),
+ expected=ragged_factory_ops.constant_value(
+ [[b'A', b'b', b'C'], [b'd', b'E']])),
dict( # shape=[D1, (D2)]
- condition=ragged.constant_value([True, False]),
- x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E']]),
- y=ragged.constant_value([['a', 'b', 'c'], ['d', 'e']]),
- expected=ragged.constant_value([[b'A', b'B', b'C'], [b'd', b'e']])),
+ condition=ragged_factory_ops.constant_value([True, False]),
+ x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+ y=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e']]),
+ expected=ragged_factory_ops.constant_value(
+ [[b'A', b'B', b'C'], [b'd', b'e']])),
#=========================================================================
# Coordinate-retrieval mode
#=========================================================================
@@ -54,24 +64,25 @@
condition=[[True, False], [False, True]],
expected=[[0, 0], [1, 1]]),
dict( # shape=[D1, (D2)]
- condition=ragged.constant_value([[True, False, True], [False, True]]),
+ condition=ragged_factory_ops.constant_value(
+ [[True, False, True], [False, True]]),
expected=[[0, 0], [0, 2], [1, 1]]),
dict( # shape=[D1, (D2), (D3)]
- condition=ragged.constant_value([
+ condition=ragged_factory_ops.constant_value([
[[True, False, True], [False, True]],
[[True], [], [False], [False, True, False]]
]),
expected=[[0, 0, 0], [0, 0, 2], [0, 1, 1],
[1, 0, 0], [1, 3, 1]]),
dict( # shape=[D1, (D2), D3]
- condition=ragged.constant_value([
+ condition=ragged_factory_ops.constant_value([
[[True, False], [False, True]],
[[True, False], [False, False], [True, False], [False, True]]
], ragged_rank=1),
expected=[[0, 0, 0], [0, 1, 1],
[1, 0, 0], [1, 2, 0], [1, 3, 1]]),
dict( # shape=[D1, (D2), (D3), (D4)]
- condition=ragged.constant_value([
+ condition=ragged_factory_ops.constant_value([
[[[], [True]]],
[[[True, False, True], [False, True]],
[[True], [], [False], [False, True, False]]]
@@ -98,44 +109,46 @@
y=[['a', 'b'], ['d', 'e']],
expected=[[b'A', b'b'], [b'd', b'E']]),
dict( # shape=[D1, (D2)]
- condition=ragged.constant_value([[True, False, True], [False, True]]),
- x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E']]),
- y=ragged.constant_value([['a', 'b', 'c'], ['d', 'e']]),
- expected=ragged.constant_value([[b'A', b'b', b'C'], [b'd', b'E']])),
+ condition=ragged_factory_ops.constant_value(
+ [[True, False, True], [False, True]]),
+ x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+ y=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e']]),
+ expected=ragged_factory_ops.constant_value(
+ [[b'A', b'b', b'C'], [b'd', b'E']])),
dict( # shape=[D1, (D2), D3]
- condition=ragged.constant_value([
+ condition=ragged_factory_ops.constant_value([
[[True, False], [False, True]],
[[True, False], [False, False], [True, False], [False, True]]
], ragged_rank=1),
- x=ragged.constant_value([
+ x=ragged_factory_ops.constant_value([
[['A', 'B'], ['C', 'D']],
[['E', 'F'], ['G', 'H'], ['I', 'J'], ['K', 'L']]
], ragged_rank=1),
- y=ragged.constant_value([
+ y=ragged_factory_ops.constant_value([
[['a', 'b'], ['c', 'd']],
[['e', 'f'], ['g', 'h'], ['i', 'j'], ['k', 'l']]
], ragged_rank=1),
- expected=ragged.constant_value([
+ expected=ragged_factory_ops.constant_value([
[[b'A', b'b'], [b'c', b'D']],
[[b'E', b'f'], [b'g', b'h'], [b'I', b'j'], [b'k', b'L']]
], ragged_rank=1)),
dict( # shape=[D1, (D2), (D3), (D4)]
- condition=ragged.constant_value([
+ condition=ragged_factory_ops.constant_value([
[[[], [True]]],
[[[True, False, True], [False, True]],
[[True], [], [False], [False, True, False]]]
]),
- x=ragged.constant_value([
+ x=ragged_factory_ops.constant_value([
[[[], ['A']]],
[[['B', 'C', 'D'], ['E', 'F']],
[['G'], [], ['H'], ['I', 'J', 'K']]]
]),
- y=ragged.constant_value([
+ y=ragged_factory_ops.constant_value([
[[[], ['a']]],
[[['b', 'c', 'd'], ['e', 'f']],
[['g'], [], ['h'], ['i', 'j', 'k']]]
]),
- expected=ragged.constant_value([
+ expected=ragged_factory_ops.constant_value([
[[[], [b'A']]],
[[[b'B', b'c', b'D'], [b'e', b'F']],
[[b'G'], [], [b'h'], [b'i', b'J', b'k']]]
@@ -151,32 +164,26 @@
expected=[[b'A', b'B'], [b'c', b'd'], [b'E', b'F']]),
dict( # shape=[D1, (D2)]
condition=[True, False, True],
- x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E'], ['F', 'G']]),
- y=ragged.constant_value([['a', 'b'], ['c'], ['d', 'e']]),
- expected=ragged.constant_value([[b'A', b'B', b'C'], [b'c'],
- [b'F', b'G']])),
+ x=ragged_factory_ops.constant_value(
+ [['A', 'B', 'C'], ['D', 'E'], ['F', 'G']]),
+ y=ragged_factory_ops.constant_value(
+ [['a', 'b'], ['c'], ['d', 'e']]),
+ expected=ragged_factory_ops.constant_value(
+ [[b'A', b'B', b'C'], [b'c'], [b'F', b'G']])),
dict( # shape=[D1, (D2), (D3), (D4)]
- condition=ragged.constant_value([True, False]),
- x=ragged.constant_value([
+ condition=ragged_factory_ops.constant_value([True, False]),
+ x=ragged_factory_ops.constant_value([
[[[], ['A']]],
[[['B', 'C', 'D'], ['E', 'F']],
[['G'], [], ['H'], ['I', 'J', 'K']]]
]),
- y=ragged.constant_value([[[['a']]], [[['b']]]]),
- expected=ragged.constant_value([[[[], [b'A']]], [[[b'b']]]])),
+ y=ragged_factory_ops.constant_value([[[['a']]], [[['b']]]]),
+ expected=ragged_factory_ops.constant_value(
+ [[[[], [b'A']]], [[[b'b']]]])),
]) # pyformat: disable
- @test_util.run_deprecated_v1
def testRaggedWhere(self, condition, expected, x=None, y=None):
- result = ragged.where(condition, x, y)
- self.assertEqual(
- getattr(result, 'ragged_rank', 0), getattr(expected, 'ragged_rank', 0))
- with self.test_session():
- result_value = self.evaluate(result)
- if hasattr(result_value, 'tolist'):
- result_value = result_value.tolist()
- if hasattr(expected, 'tolist'):
- expected = expected.tolist()
- self.assertEqual(result_value, expected)
+ result = ragged_array_ops.where(condition, x, y)
+ self.assertRaggedEqual(result, expected)
@parameterized.parameters([
dict(
@@ -185,15 +192,16 @@
error=ValueError,
message='x and y must be either both None or both non-None'),
dict(
- condition=ragged.constant_value([[True, False, True], [False, True]]),
- x=ragged.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+ condition=ragged_factory_ops.constant_value([[True, False, True],
+ [False, True]]),
+ x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
y=[['a', 'b'], ['d', 'e']],
error=ValueError,
message='Input shapes do not match.'),
])
def testRaggedWhereErrors(self, condition, error, message, x=None, y=None):
with self.assertRaisesRegexp(error, message):
- ragged.where(condition, x, y)
+ ragged_array_ops.where(condition, x, y)
if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/segment_id_ops.py b/tensorflow/python/ops/ragged/segment_id_ops.py
index fa2970c..42dc132 100644
--- a/tensorflow/python/ops/ragged/segment_id_ops.py
+++ b/tensorflow/python/ops/ragged/segment_id_ops.py
@@ -25,12 +25,14 @@
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.util.tf_export import tf_export
# For background on "segments" and "segment ids", see:
# https://www.tensorflow.org/api_guides/python/math_ops#Segmentation
+@tf_export("ragged.row_splits_to_segment_ids")
def row_splits_to_segment_ids(splits, name=None):
- """Generates the segmentation corresponding to a RaggedTensor `splits` vector.
+ """Generates the segmentation corresponding to a RaggedTensor `row_splits`.
Returns an integer vector `segment_ids`, where `segment_ids[i] == j` if
`splits[j] <= i < splits[j+1]`. Example:
@@ -63,8 +65,9 @@
# For background on "segments" and "segment ids", see:
# https://www.tensorflow.org/api_guides/python/math_ops#Segmentation
+@tf_export("ragged.segment_ids_to_row_splits")
def segment_ids_to_row_splits(segment_ids, num_segments=None, name=None):
- """Generates the RaggedTensor `splits` vector corresponding to a segmentation.
+ """Generates the RaggedTensor `row_splits` corresponding to a segmentation.
Returns an integer vector `splits`, where `splits[0] = 0` and
`splits[i] = splits[i-1] + count(segment_ids==i)`. Example:
diff --git a/tensorflow/python/ops/resource_variable_ops.py b/tensorflow/python/ops/resource_variable_ops.py
index 1066b35..dc53fb8 100644
--- a/tensorflow/python/ops/resource_variable_ops.py
+++ b/tensorflow/python/ops/resource_variable_ops.py
@@ -43,6 +43,7 @@
# pylint: enable=wildcard-import
from tensorflow.python.training.checkpointable import base as checkpointable
from tensorflow.python.util import compat
+from tensorflow.python.util.deprecation import deprecated
def get_resource_handle_data(graph_op):
@@ -685,6 +686,7 @@
raise NotImplementedError(
"numpy() is only available when eager execution is enabled.")
+ @deprecated(None, "Prefer Dataset.range instead.")
def count_up_to(self, limit):
"""Increments this variable until it reaches `limit`.
diff --git a/tensorflow/python/ops/sparse_ops.py b/tensorflow/python/ops/sparse_ops.py
index 346ab9c..097b485 100644
--- a/tensorflow/python/ops/sparse_ops.py
+++ b/tensorflow/python/ops/sparse_ops.py
@@ -299,7 +299,7 @@
@tf_export("sparse.concat", v1=[])
-def sparse_concat_v2(axis, sp_inputs, expand_nonconcat_dim=False, name=None): # pylint: disable=missing-docstring
+def sparse_concat_v2(axis, sp_inputs, expand_nonconcat_dims=False, name=None): # pylint: disable=missing-docstring
sp_inputs = _convert_to_sparse_tensors(sp_inputs)
if len(sp_inputs) == 1: # Degenerate case of one tensor.
@@ -309,7 +309,7 @@
vals = [sp_input.values for sp_input in sp_inputs]
shapes = [sp_input.dense_shape for sp_input in sp_inputs]
- if expand_nonconcat_dim:
+ if expand_nonconcat_dims:
max_shape = math_ops.reduce_max(
array_ops.concat(
[array_ops.reshape(shape, [1, -1]) for shape in shapes], 0), 0)
@@ -1093,6 +1093,9 @@
@deprecation.deprecated_endpoints("sparse_reduce_max")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
+@deprecation.deprecated_args(
+ None, "reduction_axes is deprecated, use axis instead",
+ "reduction_axes")
def sparse_reduce_max(sp_input, axis=None, keepdims=None,
reduction_axes=None, keep_dims=None):
"""Computes the max of elements across dimensions of a SparseTensor.
@@ -1141,7 +1144,7 @@
axis: The dimensions to reduce; list or scalar. If `None` (the
default), reduces all dimensions.
keepdims: If true, retain reduced dimensions with length 1.
- reduction_axes: Deprecated name of axis.
+ reduction_axes: Deprecated name of `axis`.
keep_dims: Deprecated alias for `keepdims`.
Returns:
@@ -1279,6 +1282,9 @@
@deprecation.deprecated_endpoints("sparse_reduce_sum")
@deprecation.deprecated_args(
None, "keep_dims is deprecated, use keepdims instead", "keep_dims")
+@deprecation.deprecated_args(
+ None, "reduction_axes is deprecated, use axis instead",
+ "reduction_axes")
def sparse_reduce_sum(sp_input, axis=None, keepdims=None,
reduction_axes=None, keep_dims=None):
"""Computes the sum of elements across dimensions of a SparseTensor.
@@ -1314,7 +1320,7 @@
axis: The dimensions to reduce; list or scalar. If `None` (the
default), reduces all dimensions.
keepdims: If true, retain reduced dimensions with length 1.
- reduction_axes: Deprecated name of axis.
+ reduction_axes: Deprecated name of `axis`.
keep_dims: Deprecated alias for `keepdims`.
Returns:
@@ -2696,6 +2702,7 @@
if args:
x, args = args[0], args[1:]
else:
+ kwargs = kwargs.copy()
x = kwargs.pop(self._x, None)
if isinstance(x, sparse_tensor.SparseTensor):
return sparse_tensor.SparseTensor(
diff --git a/tensorflow/python/ops/standard_ops.py b/tensorflow/python/ops/standard_ops.py
index c614d07..ba3bd09 100644
--- a/tensorflow/python/ops/standard_ops.py
+++ b/tensorflow/python/ops/standard_ops.py
@@ -22,6 +22,8 @@
import sys as _sys
+from tensorflow.python import autograph
+
# pylint: disable=g-bad-import-order
# Imports the following modules so that @RegisterGradient get executed.
from tensorflow.python.ops import array_grad
@@ -69,6 +71,8 @@
from tensorflow.python.ops.numerics import *
from tensorflow.python.ops.parsing_ops import *
from tensorflow.python.ops.partitioned_variables import *
+from tensorflow.python.ops.ragged import ragged_dispatch as _ragged_dispatch
+from tensorflow.python.ops.ragged import ragged_operators as _ragged_operators
from tensorflow.python.ops.random_ops import *
from tensorflow.python.ops.script_ops import py_func
from tensorflow.python.ops.session_ops import *
@@ -100,3 +104,7 @@
from tensorflow.python.ops.variables import *
# pylint: enable=wildcard-import
# pylint: enable=g-bad-import-order
+
+
+# These modules were imported to set up RaggedTensor operators and dispatchers:
+del _ragged_dispatch, _ragged_operators
diff --git a/tensorflow/python/ops/state_ops.py b/tensorflow/python/ops/state_ops.py
index 76684f8..71aacee 100644
--- a/tensorflow/python/ops/state_ops.py
+++ b/tensorflow/python/ops/state_ops.py
@@ -32,6 +32,8 @@
# pylint: disable=wildcard-import
from tensorflow.python.ops.gen_state_ops import *
# pylint: enable=wildcard-import
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.deprecation import deprecated
from tensorflow.python.util.tf_export import tf_export
@@ -223,6 +225,7 @@
@tf_export(v1=["count_up_to"])
+@deprecated(None, "Prefer Dataset.range instead.")
def count_up_to(ref, limit, name=None):
r"""Increments 'ref' until it reaches 'limit'.
@@ -595,7 +598,9 @@
name=name))
-@tf_export("batch_scatter_update")
+@tf_export(v1=["batch_scatter_update"])
+@deprecation.deprecated(
+ "2018-11-29", "Use the batch_scatter_update method of Variable instead.")
def batch_scatter_update(ref, indices, updates, use_locking=True, name=None):
"""Generalization of `tf.scatter_update` to axis different than 0.
diff --git a/tensorflow/python/ops/string_ops.py b/tensorflow/python/ops/string_ops.py
index b6b329c..9967f48 100644
--- a/tensorflow/python/ops/string_ops.py
+++ b/tensorflow/python/ops/string_ops.py
@@ -38,6 +38,7 @@
from tensorflow.python.ops.gen_string_ops import *
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.tf_export import tf_export
# pylint: enable=g-bad-import-order
# pylint: enable=wildcard-import
@@ -45,6 +46,7 @@
# pylint: disable=redefined-builtin
@tf_export("strings.regex_full_match")
+@dispatch.add_dispatch_support
def regex_full_match(input, pattern, name=None):
r"""Match elements of `input` with regex `pattern`.
@@ -76,6 +78,7 @@
@tf_export(
"strings.regex_replace", v1=["strings.regex_replace", "regex_replace"])
@deprecation.deprecated_endpoints("regex_replace")
+@dispatch.add_dispatch_support
def regex_replace(input, pattern, rewrite, replace_global=True, name=None):
r"""Replace elements of `input` matching regex `pattern` with `rewrite`.
@@ -350,10 +353,13 @@
# This wrapper provides backwards compatibility for code that predates the
# unit argument and that passed 'name' as a positional argument.
@tf_export(v1=["strings.length"])
+@dispatch.add_dispatch_support
def string_length(input, name=None, unit="BYTE"):
return gen_string_ops.string_length(input, unit=unit, name=name)
+
@tf_export("strings.length", v1=[])
+@dispatch.add_dispatch_support
def string_length_v2(input, unit="BYTE", name=None):
return string_length(input, name, unit)
@@ -361,7 +367,7 @@
string_length.__doc__ = gen_string_ops.string_length.__doc__
-@tf_export("substr")
+@tf_export(v1=["substr"])
@deprecation.deprecated(None, "Use `tf.strings.substr` instead of `tf.substr`.")
def substr_deprecated(input, pos, len, name=None, unit="BYTE"):
return substr(input, pos, len, name=name, unit=unit)
@@ -370,16 +376,19 @@
@tf_export(v1=["strings.substr"])
+@dispatch.add_dispatch_support
def substr(input, pos, len, name=None, unit="BYTE"):
return gen_string_ops.substr(input, pos, len, unit=unit, name=name)
+substr.__doc__ = gen_string_ops.substr.__doc__
+
@tf_export("strings.substr", v1=[])
+@dispatch.add_dispatch_support
def substr_v2(input, pos, len, unit="BYTE", name=None):
- return substr(input, pos, len, name=name, unit=unit)
+ return gen_string_ops.substr(input, pos, len, unit=unit, name=name)
-
-substr.__doc__ = gen_string_ops.substr.__doc__
+substr_v2.__doc__ = gen_string_ops.substr.__doc__
ops.NotDifferentiable("RegexReplace")
@@ -395,6 +404,7 @@
@tf_export("strings.to_number", v1=[])
+@dispatch.add_dispatch_support
def string_to_number(input, out_type=dtypes.float32, name=None):
r"""Converts each string in the input Tensor to the specified numeric type.
@@ -418,6 +428,7 @@
@tf_export("strings.to_hash_bucket", v1=[])
+@dispatch.add_dispatch_support
def string_to_hash_bucket(input, num_buckets, name=None):
# pylint: disable=line-too-long
r"""Converts each string in the input Tensor to its hash mod by a number of buckets.
diff --git a/tensorflow/python/ops/summary_op_util.py b/tensorflow/python/ops/summary_op_util.py
index c72a9ae..93d8d50 100644
--- a/tensorflow/python/ops/summary_op_util.py
+++ b/tensorflow/python/ops/summary_op_util.py
@@ -21,10 +21,10 @@
import contextlib
import re
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.platform import tf_logging
-from tensorflow.python.training import distribution_strategy_context
def collect(val, collections, default_collections):
diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py
index d151694..85333ee 100644
--- a/tensorflow/python/ops/tensor_array_ops.py
+++ b/tensorflow/python/ops/tensor_array_ops.py
@@ -20,10 +20,8 @@
from __future__ import print_function
import contextlib
-import os
import weakref
-from tensorflow.python import tf2
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -32,6 +30,7 @@
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import gen_control_flow_ops
from tensorflow.python.ops import gen_data_flow_ops
from tensorflow.python.ops import list_ops
@@ -40,10 +39,6 @@
from tensorflow.python.util.tf_export import tf_export
-ENABLE_TENSOR_ARRAY_V2 = (
- tf2.enabled() or os.getenv("TF_ENABLE_TENSOR_ARRAY_V2") is not None)
-
-
# _GraphTensorArray accesses many of the hidden generated ops, but is in
# fact built to wrap these methods.
# pylint: disable=protected-access
@@ -1013,7 +1008,7 @@
if context.executing_eagerly():
implementation = _EagerTensorArray
else:
- if ENABLE_TENSOR_ARRAY_V2:
+ if control_flow_util.ENABLE_CONTROL_FLOW_V2:
implementation = _GraphTensorArrayV2
else:
implementation = _GraphTensorArray
diff --git a/tensorflow/python/ops/variables.py b/tensorflow/python/ops/variables.py
index 4824c92..e231343 100644
--- a/tensorflow/python/ops/variables.py
+++ b/tensorflow/python/ops/variables.py
@@ -637,37 +637,84 @@
"""
raise NotImplementedError
+ def batch_scatter_update(self, sparse_delta, use_locking=False, name=None):
+ """Assigns `IndexedSlices` to this variable batch-wise.
+
+ Analogous to `batch_gather`. This assumes that this variable and the
+ sparse_delta IndexedSlices have a series of leading dimensions that are the
+ same for all of them, and the updates are performed on the last dimension of
+ indices. In other words, the dimensions should be the following:
+
+ `num_prefix_dims = sparse_delta.indices.ndims - 1`
+ `batch_dim = num_prefix_dims + 1`
+ `sparse_delta.updates.shape = sparse_delta.indices.shape + var.shape[
+ batch_dim:]`
+
+ where
+
+ `sparse_delta.updates.shape[:num_prefix_dims]`
+ `== sparse_delta.indices.shape[:num_prefix_dims]`
+ `== var.shape[:num_prefix_dims]`
+
+ And the operation performed can be expressed as:
+
+ `var[i_1, ..., i_n,
+ sparse_delta.indices[i_1, ..., i_n, j]] = sparse_delta.updates[
+ i_1, ..., i_n, j]`
+
+ When sparse_delta.indices is a 1D tensor, this operation is equivalent to
+ `scatter_update`.
+
+ To avoid this operation one can looping over the first `ndims` of the
+ variable and using `scatter_update` on the subtensors that result of slicing
+ the first dimension. This is a valid option for `ndims = 1`, but less
+ efficient than this implementation.
+
+ Args:
+ sparse_delta: `IndexedSlices` to be assigned to this variable.
+ use_locking: If `True`, use locking during the operation.
+ name: the name of the operation.
+
+ Returns:
+ A `Tensor` that will hold the new value of this variable after
+ the scattered subtraction has completed.
+
+ Raises:
+ ValueError: if `sparse_delta` is not an `IndexedSlices`.
+ """
+ raise NotImplementedError
+
def scatter_nd_sub(self, indices, updates, name=None):
"""Applies sparse subtraction to individual values or slices in a Variable.
- `ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
+ Assuming the variable has rank `P` and `indices` is a `Tensor` of rank `Q`.
- `indices` must be integer tensor, containing indices into `ref`.
+ `indices` must be integer tensor, containing indices into self.
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 slices (if `K < P`) along the `K`th
- dimension of `ref`.
+ dimension of self.
`updates` is `Tensor` of rank `Q-1+P-K` with shape:
```
- [d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
+ [d_0, ..., d_{Q-2}, self.shape[K], ..., self.shape[P-1]].
```
For example, say we want to add 4 scattered elements to a rank-1 tensor to
8 elements. In Python, that update would look like this:
```python
- ref = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
+ v = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
indices = tf.constant([[4], [3], [1] ,[7]])
updates = tf.constant([9, 10, 11, 12])
- op = ref.scatter_nd_sub(indices, updates)
+ op = v.scatter_nd_sub(indices, updates)
with tf.Session() as sess:
print sess.run(op)
```
- The resulting update to ref would look like this:
+ The resulting update to v would look like this:
[1, -9, 3, -6, -6, 6, 7, -4]
@@ -691,34 +738,34 @@
def scatter_nd_add(self, indices, updates, name=None):
"""Applies sparse addition to individual values or slices in a Variable.
- `ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
+ The Variable has rank `P` and `indices` is a `Tensor` of rank `Q`.
- `indices` must be integer tensor, containing indices into `ref`.
+ `indices` must be integer tensor, containing indices into self.
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 slices (if `K < P`) along the `K`th
- dimension of `ref`.
+ dimension of self.
`updates` is `Tensor` of rank `Q-1+P-K` with shape:
```
- [d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
+ [d_0, ..., d_{Q-2}, self.shape[K], ..., self.shape[P-1]].
```
For example, say we want to add 4 scattered elements to a rank-1 tensor to
8 elements. In Python, that update would look like this:
```python
- ref = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
+ v = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
indices = tf.constant([[4], [3], [1] ,[7]])
updates = tf.constant([9, 10, 11, 12])
- add = ref.scatter_nd_add(indices, updates)
+ add = v.scatter_nd_add(indices, updates)
with tf.Session() as sess:
print sess.run(add)
```
- The resulting update to ref would look like this:
+ The resulting update to v would look like this:
[1, 13, 3, 14, 14, 6, 7, 20]
@@ -742,34 +789,34 @@
def scatter_nd_update(self, indices, updates, name=None):
"""Applies sparse assignment to individual values or slices in a Variable.
- `ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
+ The Variable has rank `P` and `indices` is a `Tensor` of rank `Q`.
- `indices` must be integer tensor, containing indices into `ref`.
+ `indices` must be integer tensor, containing indices into self.
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 slices (if `K < P`) along the `K`th
- dimension of `ref`.
+ dimension of self.
`updates` is `Tensor` of rank `Q-1+P-K` with shape:
```
- [d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
+ [d_0, ..., d_{Q-2}, self.shape[K], ..., self.shape[P-1]].
```
For example, say we want to add 4 scattered elements to a rank-1 tensor to
8 elements. In Python, that update would look like this:
```python
- ref = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
+ v = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
indices = tf.constant([[4], [3], [1] ,[7]])
updates = tf.constant([9, 10, 11, 12])
- op = ref.scatter_nd_assign(indices, updates)
+ op = v.scatter_nd_assign(indices, updates)
with tf.Session() as sess:
print sess.run(op)
```
- The resulting update to ref would look like this:
+ The resulting update to v would look like this:
[1, 11, 3, 10, 9, 6, 7, 12]
@@ -790,6 +837,7 @@
"""
raise NotImplementedError
+ @deprecated(None, "Prefer Dataset.range instead.")
def count_up_to(self, limit):
"""Increments this variable until it reaches `limit`.
@@ -1842,6 +1890,55 @@
use_locking=use_locking,
name=name)
+ def batch_scatter_update(self, sparse_delta, use_locking=False, name=None):
+ """Assigns `IndexedSlices` to this variable batch-wise.
+
+ Analogous to `batch_gather`. This assumes that this variable and the
+ sparse_delta IndexedSlices have a series of leading dimensions that are the
+ same for all of them, and the updates are performed on the last dimension of
+ indices. In other words, the dimensions should be the following:
+
+ `num_prefix_dims = sparse_delta.indices.ndims - 1`
+ `batch_dim = num_prefix_dims + 1`
+ `sparse_delta.updates.shape = sparse_delta.indices.shape + var.shape[
+ batch_dim:]`
+
+ where
+
+ `sparse_delta.updates.shape[:num_prefix_dims]`
+ `== sparse_delta.indices.shape[:num_prefix_dims]`
+ `== var.shape[:num_prefix_dims]`
+
+ And the operation performed can be expressed as:
+
+ `var[i_1, ..., i_n,
+ sparse_delta.indices[i_1, ..., i_n, j]] = sparse_delta.updates[
+ i_1, ..., i_n, j]`
+
+ When sparse_delta.indices is a 1D tensor, this operation is equivalent to
+ `scatter_update`.
+
+ To avoid this operation one can looping over the first `ndims` of the
+ variable and using `scatter_update` on the subtensors that result of slicing
+ the first dimension. This is a valid option for `ndims = 1`, but less
+ efficient than this implementation.
+
+ Args:
+ sparse_delta: `IndexedSlices` to be assigned to this variable.
+ use_locking: If `True`, use locking during the operation.
+ name: the name of the operation.
+
+ Returns:
+ A `Tensor` that will hold the new value of this variable after
+ the scattered subtraction has completed.
+
+ Raises:
+ ValueError: if `sparse_delta` is not an `IndexedSlices`.
+ """
+ return state_ops.batch_scatter_update(
+ self, sparse_delta.indices, sparse_delta.values,
+ use_locking=use_locking, name=name)
+
def scatter_nd_sub(self, indices, updates, name=None):
"""Applies sparse subtraction to individual values or slices in a Variable.
@@ -2021,6 +2118,7 @@
new_axis_mask=new_axis_mask,
shrink_axis_mask=shrink_axis_mask)
+ @deprecated(None, "Prefer Dataset.range instead.")
def count_up_to(self, limit):
"""Increments this variable until it reaches `limit`.
diff --git a/tensorflow/python/ops/while_v2.py b/tensorflow/python/ops/while_v2.py
index 59ca29e..3e5a8fc 100644
--- a/tensorflow/python/ops/while_v2.py
+++ b/tensorflow/python/ops/while_v2.py
@@ -23,6 +23,7 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import func_graph as func_graph_module
@@ -51,13 +52,6 @@
# to them and then pass those in as data inputs. This should probably be
# handled in the CapturingGraph itself.
-# Op types that output a resource tensor representing a TensorArray handle.
-TENSOR_ARRAY_HANDLE_OPS = (
- "TensorArrayV3",
- "TensorArrayGradV3",
- "TensorArrayGradWithShape",
-)
-
def while_loop(cond,
body,
@@ -105,7 +99,7 @@
# Automatic control dependencies are added in defuns, but not in v1
# graphs. Propagate that behavior here.
- add_control_dependencies = util.in_defun()
+ add_control_dependencies = ops.get_default_graph()._add_control_dependencies
# Build a `cond` wrapper that can handle the extra counter loop_var.
def wrapped_cond(loop_counter, *args):
@@ -200,30 +194,6 @@
" this as a loop variable." % str(external_capture))
cond_graph.capture(external_capture)
- # Export all tensors in the loop body that may be needed for gradient
- # computation. We do this by accumulating the intermediate values in
- # TensorLists.
- intermediate_tensors = _get_intermediates(body_graph)
-
- for intermediate_tensor in intermediate_tensors:
- tensor_list = list_ops.empty_tensor_list(
- element_dtype=intermediate_tensor.dtype,
- element_shape=intermediate_tensor.shape,
- max_num_elements=maximum_iterations)
- loop_vars.append(tensor_list)
- with cond_graph.as_default():
- # Add a placeholder to cond_graph's inputs corresponding to the
- # tensor_list.
- cond_graph.capture(tensor_list)
- with body_graph.as_default():
- # Push the intermediate tensor to the tensor list. This captures the
- # `tensor_list` as well.
- appended_tensor_list = list_ops.tensor_list_push_back(
- tensor_list,
- intermediate_tensor)
- # Add this modified tensor list to the list of outputs.
- body_graph.outputs.append(appended_tensor_list)
-
# Make sure that the shapes of the loop outputs are compatible with the
# shape invariants, or the shapes of the loop vars if the invariants are not
# specified.
@@ -272,33 +242,32 @@
@ops.RegisterGradient("While")
def _WhileGrad(op, *grads): # pylint: disable=invalid-name
"""The gradient of a While op produced by while_loop."""
- body_graph = _get_body_graph(op)
+ cond_graph = _get_graph(op, "cond")
+ body_graph = _get_graph(op, "body")
+ orig_num_params = len(body_graph.outputs)
- # Set the incoming gradient of TensorArray handles to None. The gradient
- # implementation currently assumes all resource tensors correspond to float32
- # ResourceVariables, which can lead to runtime shape errors when used with a
- # TensorArray. This is a workaround until TensorArrays are reimplemented with
- # TensorLists instead of resources.
- # Also set the incoming gradient of non-trainable inputs to None. It is
- # possible that we receive non-None gradients for non-trainable types in
- # nested while loops because we accumulate outputs of the inner while as
- # variant tensors which are trainable and hence receive zeros_like tensors in
- # the gradient pass. The non-trainable tensors then receive the popped zeros
- # tensor from this zeros variant. The gradient for the loop vars corresponding
- # to these tensors is None or zeros (this happens only if the loop var is
- # accumulated as well) in _grad_fn so we reset these.
+ maximum_iterations = op.get_attr(
+ "_maximum_iterations") if _is_in_xla_context() else None
+ assert not _is_in_xla_context() or maximum_iterations is not None
+
+ # Set the incoming gradient of non-trainable inputs to None. It is possible
+ # that we receive non-None gradients for non-trainable types in nested while
+ # loops because we accumulate outputs of the inner while as variant tensors
+ # which are trainable and hence receive zeros_like tensors in the gradient
+ # pass. The non-trainable tensors then receive the popped zeros tensor from
+ # this zeros variant. The gradient for the loop vars corresponding to these
+ # tensors is None or zeros (this happens only if the loop var is accumulated
+ # as well) in _grad_fn so we reset these.
# TODO(b/118712257): Remove the IsTrainable filter once we can handle None
# output grads in _grad_fn.
grads = [
- None if _is_tensor_array_handle(output) or
- not gradients_impl.IsTrainable(output) else grad
- for grad, output in zip(grads, op.outputs)
+ None if not _is_trainable(output) else grad
+ for grad, output in zip(grads, body_graph.outputs)
]
# Ensure that all non-resource trainable outputs have incoming gradients.
- assert all(g is not None or o.dtype == dtypes.resource or
- not gradients_impl.IsTrainable(o)
- for o, g in zip(op.outputs, grads)
+ assert all(g is not None or o.dtype == dtypes.resource or not _is_trainable(o)
+ for o, g in zip(body_graph.outputs, grads)
), "All trainable loop vars must receive incoming gradients."
# We compute the gradient for the sub-graph between trainable ys and xs
# with non-None incoming gradients. We later pad the None's to the list of
@@ -307,32 +276,36 @@
body_graph.outputs, body_graph.inputs, grads) if grad is not None])
body_grad_graph, args = _create_grad_func(
- ys, xs, non_none_grads, body_graph,
- util.unique_grad_fn_name(body_graph.name), op)
+ ys, xs, non_none_grads, cond_graph, body_graph,
+ util.unique_grad_fn_name(body_graph.name), op, maximum_iterations)
- intermediate_tensors = _get_intermediates(body_grad_graph)
+ if body_grad_graph.while_op_needs_rewrite:
+ # Modify 'op' to output the intermediate accumulators needed by the grad
+ # function.
+ # NOTE(skyewm): if there are any active sessions, this modification to `op`
+ # may make them unrunnable!
- maximum_iterations = op.get_attr(
- "_maximum_iterations") if _is_in_xla_context() else None
- assert not _is_in_xla_context() or maximum_iterations is not None
- for intermediate_tensor in intermediate_tensors:
- tensor_list = list_ops.empty_tensor_list(
- element_dtype=intermediate_tensor.dtype,
- element_shape=intermediate_tensor.shape,
- max_num_elements=maximum_iterations)
+ cond_graph.name += "_rewritten"
+ body_graph.name += "_rewritten"
- with body_grad_graph.as_default():
- tensor_list_ph = body_grad_graph.capture(tensor_list, whitelisted=True)
- # Push the intermediate tensor to the tensor list.
- appended_tensor_list = list_ops.tensor_list_push_back(tensor_list_ph,
- intermediate_tensor)
- # Add this modified tensor list to the list of outputs.
- body_grad_graph.outputs.append(appended_tensor_list)
+ new_inputs = body_grad_graph.empty_tensor_lists
+ new_outputs = body_graph.outputs[orig_num_params:]
+
+ op._set_func_attr("cond", util.create_new_tf_function(cond_graph))
+ op._set_func_attr("body", util.create_new_tf_function(body_graph))
+ op._set_type_list_attr("T", body_graph.output_types)
+ op._set_shape_list_attr("output_shapes", body_graph.output_shapes)
+ op._add_while_inputs(new_inputs)
+ op._add_outputs([t.dtype for t in new_outputs],
+ [t.shape for t in new_outputs])
+ _copy_handle_data(new_outputs, op.outputs[orig_num_params:])
+
+ captured_inputs = _resolve_grad_captures(body_graph, body_grad_graph, op)
+ loop_vars = args + captured_inputs
def grad_cond(counter, max_iters, *unused_args):
return counter < max_iters
- loop_vars = args + body_grad_graph.external_captures
grad_cond_name = util.unique_grad_fn_name(op.get_attr("cond").name)
cond_grad_graph = func_graph_module.func_graph_from_py_func(
grad_cond_name, grad_cond, loop_vars, {},
@@ -354,8 +327,7 @@
# See comment in while_loop.
outputs = [array_ops.identity(t) for t in outputs]
- # Set None as the output gradient for tensors with None input gradient
- # e.g. TensorArray handles.
+ # Set None as the output gradient for tensors with None input gradient.
# outputs[0] is the loop counter.
# outputs[1] is the total number of loop iterations.
index = 2
@@ -369,6 +341,24 @@
return none_padded_outputs
+def _is_trainable(tensor):
+ """Returns whether the given tensor is trainable."""
+ if not gradients_impl.IsTrainable(tensor):
+ return False
+
+ # Special case: untrainable accumulator output. The gradients algorithm
+ # doesn't know about tensor lists of untrainable elements. In theory the
+ # tensor list gradient functions should return None as appropriate, but
+ # because we can't return None from the gradient function we filter out
+ # untrainable accumulator output here to avoid computing the gradient at all.
+ if tensor.op.type == "TensorListPopBack" and tensor.value_index == 0:
+ assert tensor.dtype == dtypes.variant
+ element_type = tensor.op.get_attr("element_dtype")
+ return gradients_impl.IsTrainable(element_type)
+
+ return True
+
+
def _validate_and_convert_to_tensor(maximum_iterations):
"""Checks that `maximum_iterations` is valid.
@@ -410,20 +400,21 @@
# TODO(srbs): Pull this into common utils for cond_v2 and while_v2.
-def _get_body_graph(while_op):
- """Returns `FuncGraph` for the while body.
+def _get_graph(while_op, func_attr_name):
+ """Returns `FuncGraph` for the given function attribute.
Args:
while_op: The While Operation.
+ func_attr_name: string
Returns:
- `FuncGraph` for the while body.
+ `FuncGraph`
"""
# TODO(srbs): Handle TensorShapeProto in function_def_to_graph.input_shapes.
input_shapes = [
tensor_shape.TensorShape(s) for s in while_op.get_attr("output_shapes")
]
- func_name = while_op.get_attr("body").name
+ func_name = while_op.get_attr(func_attr_name).name
fdef = while_op.graph._get_function(func_name).definition
# `while_op.graph` may not be the same as `ops.get_default_graph()` e.g.
# if the `while_op` is in the body of another if/while/defun. We build the
@@ -436,7 +427,8 @@
return func_graph
-def _create_grad_func(ys, xs, grads, func_graph, name, while_op):
+def _create_grad_func(ys, xs, grads, cond_graph, body_graph, name, while_op,
+ max_iters):
"""Builds and returns the gradient FuncGraph of `func_graph` and its args.
The returned grad_func_graph must be called with the returned
@@ -446,9 +438,11 @@
ys: A `Tensor` or list of tensors to be differentiated.
xs: A `Tensor` or list of tensors to be used for differentiation.
grads: The incoming grads for `ys`.
- func_graph: FuncGraph for the forward body function.
+ cond_graph: FuncGraph for the forward cond function.
+ body_graph: FuncGraph for the forward body function.
name: Name of the returned gradient function.
while_op: The forward While op.
+ max_iters: the maximum number of iterations, or None if no limit.
Returns:
2-tuple of (grad_func_graph, args).
@@ -464,9 +458,10 @@
# `external_captures`.
grad_func_graph = func_graph_module.func_graph_from_py_func(
name,
- lambda *args: _grad_fn(ys, xs, args, func_graph),
+ lambda *args: _grad_fn(ys, xs, args, body_graph),
args, {},
- func_graph=_WhileBodyGradFuncGraph(name, func_graph))
+ func_graph=_WhileBodyGradFuncGraph(name, cond_graph, body_graph,
+ max_iters))
# Add the popped accumulators to the list of outputs.
for internal_capture in grad_func_graph.internal_captures:
@@ -506,10 +501,11 @@
# Build the gradient graph. Note that this builds the gradient computation of
# func_graph in the current graph, which requires capturing tensors from
# func_graph. The captured func_graph tensors are resolved to external tensors
- # in _resolve_grad_inputs.
+ # after the forward While op has been rewritten in _resolve_grad_captures.
# TODO(srbs): Mark GradientsHelper as public?
grad_outs = gradients_impl._GradientsHelper(
- ys, xs, grad_ys=grad_ys, src_graph=func_graph)
+ ys, xs, grad_ys=grad_ys, src_graph=func_graph,
+ unconnected_gradients="zero")
# TODO(b/118712257): Handle the case when grad_outs has None's e.g. when there
# is a tf.StopGradient in the loop body.
@@ -519,43 +515,45 @@
return [counter + 1, total_iters] + grad_outs
-def _get_intermediates(func_graph):
- """Returns all tensors in `func_graph` that should be accumulated."""
- # We currently accumulate output tensors of most ops in the function and rely
- # on the pruning pass to get rid of the unused accumulators at runtime.
- # However, this can bloat the GraphDef and make debugging harder so we perform
- # some optimizations.
- #
- # Optimization we currently perform:
- # 1. We do not accumulate tensors which already have an accumulator
- # in the loop body.
- # 2. We do not accumulate outputs of Identity nodes. When building the
- # FuncGraph, we add an Identity node for each output (see
- # `AutomaticControlDependencies.mark_as_return`). Accumulating outputs
- # of all these nodes bloats the GraphDef quite a bit so we remove those.
- # Since the gradient of an Identity node does not rely on its forward op's
- # input this is safe to do.
- #
- # Other possible optimizations:
- # 1. Only accumulate tensors that will be required by the backward pass.
- # This will require running the gradient pass and hence would increase the
- # graph building time for the forward pass.
- # 2. Do not accumulate Const nodes created inside the loop body.
- # 3. Do not accumulate inputs that are passed as-is, e.g. loop invariants.
- # TODO(srbs): 2 and 3 may be hard optimizations for the runtime optimizer
- # since it requires knowledge of the while loop semantics. If so, consider
- # doing those here.
- intermediates = []
+def _resolve_grad_captures(body_graph, body_grad_graph, while_op):
+ """Returns the tensors to pass as captured inputs to `body_grad_graph`.
- for op in func_graph.get_operations():
- if op.type == "Identity":
- continue
- for o in op.outputs:
- if (o != func_graph.inputs[0] and # Loop counter.
- o.dtype != dtypes.resource and # Do not accumulate resource tensors.
- _get_accumulator(o) is None): # Has existing accumulator.
- intermediates.append(o)
- return intermediates
+ `body_grad_graph` may have external references to:
+ 1. Its outer graph containing the input gradients. These are left as-is.
+ 2. Accumulators captured from the forward-pass graph. These should have been
+ added as `while_op` outputs after the gradient graph was built. We replace
+ these with the corresponding output of `while_op`, i.e. a tensor in
+ `body_graph.outer_graph`. In the case of nested control flow or functions,
+ the gradient logic handling `body_grad_graph.outer_graph` will make sure
+ the tensor from `body_graph.outer_graph` is also correctly captured.
+
+ Args:
+ body_graph: FuncGraph. The forward-pass body function.
+ body_grad_graph: FuncGraph. The body gradients function.
+ while_op: The forward-pass While Operation calling `body_graph`.
+
+ Returns:
+ A list of input tensors to be passed as the captured inputs to
+ `body_grad_graph`.
+ """
+ new_capture_inputs = []
+ for t in body_grad_graph.external_captures:
+ # All values captured by gradient computation should be from the forward
+ # graph or a captured resource variable (note that input gradients are
+ # regular non-captured inputs).
+ if t.graph == body_graph:
+ # Captured accumulator
+ t = while_op.outputs[t.graph.outputs.index(t)]
+ # Note: We rely on the capturing logic of the gradient While op graph to
+ # correctly capture the tensors in `body_graph.outer_graph`. Both cond_v2
+ # and while_v2 handle this while building their gradient functions.
+ assert t.graph == body_graph.outer_graph
+ else:
+ # Captured resource variable
+ assert t.dtype == dtypes.resource
+
+ new_capture_inputs.append(t)
+ return new_capture_inputs
def _get_accumulator(tensor):
@@ -629,9 +627,10 @@
b. Lookup the corresponding resource tensor in the forward outer graph and
try to capture that.
2. If the tensor is not of resource type:
- a. Find the accumulator for that tensor.
- b. Capture the forward While op output tensor corresponding to the
- accumulator in this FuncGraph.
+ a. Create an accumulator for that tensor and output it from the forward
+ pass. Note this also requires adding it as an input to the forward pass.
+ b. Capture the accumulator from the forward pass in this FuncGraph. This
+ will later be resolved to the correct output of the forward While op.
c. Pop a value from the captured placeholder and use it as the captured
value for the forward pass tensor.
@@ -645,16 +644,25 @@
tensor.
Attributes:
- popped_tensor_lists: Dict from the captured accumulator placeholder to the
+ while_op_needs_rewrite: True if any non-resource intermediates were
+ captured, meaning the forward While op needs to be rewritten to output the
+ corresponding accumulators.
+ empty_tensor_lists: list of EmptyTensorList tensors to be used as initial
+ input to the new accumulators in the forward graph.
+ popped_tensor_lists: dict from the captured accumulator placeholder to the
TensorList obtained after popping the intermediate tensor from it. The
values of this dict need to be added to the list of outputs.
"""
- def __init__(self, name, forward_graph):
+ def __init__(self, name, forward_cond_graph, forward_body_graph, max_iters):
super(_WhileBodyGradFuncGraph, self).__init__(name)
+ self.empty_tensor_lists = []
self.popped_tensor_lists = {}
# FuncGraph for the body of the forward While op.
- self._forward_graph = forward_graph
+ self._forward_graph = forward_body_graph
+ # FuncGraph for the cond of the forward While op.
+ self._forward_cond_graph = forward_cond_graph
+ self._maximum_iterations = max_iters
# Dict from forward intermediate tensor to its indirectly captured tensor
# in this graph. Indirect capturing happens in two ways:
# 1. For non-resource tensors we capture their accumulators from the forward
@@ -663,13 +671,10 @@
# 2. For resource tensors we directly capture their corresponding tensor
# in the forward outer graph.
self._indirect_captures = {}
- # Dict from forward graph tensor to its corresponding tensor in
- # `forward_graph.outer_graph`. For a non-resource tensor the value is the
- # forward While op's "output" corresponding its accumulator. For a resource
- # tensor it is the While op's "input" for the resource. Note: We disallow
- # creation of resources inside the while loop so if a resource tensor exists
- # inside while loop it must be a loop input.
- self._inner_to_outer_tensor = {}
+
+ @property
+ def while_op_needs_rewrite(self):
+ return self.empty_tensor_lists
def capture(self, tensor, name=None, whitelisted=False):
"""Selectively captures external tensors.
@@ -708,10 +713,6 @@
captured_tensor = self._indirect_captures.get(tensor)
if captured_tensor is not None:
- # For GradientTape housekeeping.
- assert self._inner_to_outer_tensor[tensor] in self.captures
- super(_WhileBodyGradFuncGraph, self)._capture_helper(
- self._inner_to_outer_tensor[tensor], name)
return captured_tensor
if tensor.dtype == dtypes.resource:
@@ -736,36 +737,47 @@
index], "Resource tensors must be loop invariants %s." % str(
self._forward_graph._while.inputs[index])
tensor_in_outer_graph = self._forward_graph._while.inputs[index]
- self._inner_to_outer_tensor[tensor] = tensor_in_outer_graph
self._indirect_captures[tensor] = self.capture(
tensor_in_outer_graph, whitelisted=True)
return self._indirect_captures[tensor]
- assert tensor not in self._inner_to_outer_tensor
-
- accumulator = None
-
- # Find the TensorList that was used to accumulate the tensors of this
- # intermediate tensor.
+ # Create or find an existing accumulator output for `tensor` in the forward
+ # graph, and fetch from this accumulator in the gradient graph to get the
+ # raw intermediate value.
accumulator = _get_accumulator(tensor)
if accumulator is None:
- raise ValueError("Reference to un-accumulated intermediate tensor: ",
- tensor.name)
- assert accumulator.graph == self._forward_graph
- # Get the While op output corresponding to the accumulator.
- accumulator = self._forward_graph._while.outputs[self._forward_graph.outputs
- .index(accumulator)]
+ # Create the initial empty tensor list.
+ with self._forward_graph.outer_graph.as_default():
+ tensor_list = list_ops.empty_tensor_list(
+ element_dtype=tensor.dtype, element_shape=tensor.shape,
+ max_num_elements=self._maximum_iterations)
+ self.empty_tensor_lists.append(tensor_list)
- assert accumulator.graph == self._forward_graph.outer_graph
- self._inner_to_outer_tensor[tensor] = accumulator
+ # Push the intermediate tensor to the tensor list. This captures
+ # `tensor_list`.
+ with self._forward_graph.as_default():
+ accumulator = list_ops.tensor_list_push_back(tensor_list, tensor)
+ # Add the modified tensor list to the list of outputs. This output will be
+ # all the accumulated values.
+ self._forward_graph.outputs.append(accumulator)
- # Capture the `accumulator`.
- accumulator_ph = super(_WhileBodyGradFuncGraph, self)._capture_helper(
+ # Capture in the cond graph as well so the forward cond and body inputs
+ # match.
+ with self._forward_cond_graph.as_default():
+ self._forward_cond_graph.capture(tensor_list)
+
+ # Capture the accumulator tensor list in the gradient graph directly from
+ # the forward graph -- we'll later modify this to capture the final list
+ # output by the forward While op instead.
+ captured_accumulator = super(_WhileBodyGradFuncGraph, self)._capture_helper(
accumulator, name)
+
+ # Pop the intermediate value from the tensor list in the gradient graph.
new_tensor_list, captured_tensor = list_ops.tensor_list_pop_back(
- accumulator_ph, element_dtype=tensor.dtype)
+ captured_accumulator, element_dtype=tensor.dtype)
+
self._indirect_captures[tensor] = captured_tensor
- self.popped_tensor_lists[accumulator_ph] = new_tensor_list
+ self.popped_tensor_lists[captured_accumulator] = new_tensor_list
return captured_tensor
@@ -828,28 +840,6 @@
return "Base"
-def _is_tensor_array_handle(tensor):
- """Returns whether tensor is a TensorArray handle."""
- if tensor.dtype != dtypes.resource:
- return False
-
- if tensor.op.type == "While":
- # We assume that any resource outputs of a While op correspond to a captured
- # resource input (as opposed to a loop variable specified by the user).
- # NOTE(skyewm): we could actually check this, but I can't think of when you
- # would have a resource loop variable.
- tensor = tensor.op.inputs[tensor.value_index]
-
- # TODO(b/118452219): add test coverage for this.
- tensor = func_graph_module.maybe_captured(tensor)
-
- if isinstance(tensor, ops.EagerTensor):
- # Eager execution doesn't quite support legacy tensorarray
- return False
-
- return tensor.op.type in TENSOR_ARRAY_HANDLE_OPS
-
-
def _pack_sequence_as(structure_with_tas, loop_vars):
"""Like `nest.pack_sequence_as` but also replaces flows with TensorArrays."""
diff --git a/tensorflow/python/platform/googletest.py b/tensorflow/python/platform/googletest.py
index 4d34c50..5b20e36 100644
--- a/tensorflow/python/platform/googletest.py
+++ b/tensorflow/python/platform/googletest.py
@@ -142,7 +142,7 @@
return False
-@tf_export('test.StubOutForTesting')
+@tf_export(v1=['test.StubOutForTesting'])
class StubOutForTesting(object):
"""Support class for stubbing methods out for unit testing.
diff --git a/tensorflow/python/profiler/internal/run_metadata_test.py b/tensorflow/python/profiler/internal/run_metadata_test.py
index a8859f8..f96d721 100644
--- a/tensorflow/python/profiler/internal/run_metadata_test.py
+++ b/tensorflow/python/profiler/internal/run_metadata_test.py
@@ -169,7 +169,7 @@
ret = _extract_node(run_meta, 'MatMul:MatMul')
self.assertEqual(len(ret), 0)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testLoopCPU(self):
ops.reset_default_graph()
with ops.device('/cpu:0'):
diff --git a/tensorflow/python/profiler/model_analyzer_test.py b/tensorflow/python/profiler/model_analyzer_test.py
index 8648f0b..1c7c15b 100644
--- a/tensorflow/python/profiler/model_analyzer_test.py
+++ b/tensorflow/python/profiler/model_analyzer_test.py
@@ -76,6 +76,7 @@
' ScalarW (1, 1/1 params)\n',
lib.CheckAndRemoveDoc(f.read()))
+ @test_util.run_v1_only('b/120545219')
def testSelectEverythingDetail(self):
ops.reset_default_graph()
dev = '/device:GPU:0' if test.is_gpu_available() else '/device:CPU:0'
@@ -203,6 +204,7 @@
lib.CheckAndRemoveDoc(f.read())[0:80])
# pylint: enable=line-too-long
+ @test_util.run_v1_only('b/120545219')
def testComplexCodeView(self):
ops.reset_default_graph()
outfile = os.path.join(test.get_temp_dir(), 'dump')
@@ -619,6 +621,7 @@
else:
self.assertEqual(len(gfile.ListDirectory(profile_dir)), 0)
+ @test_util.run_v1_only('b/120545219')
def testAutoProfiling(self):
ops.reset_default_graph()
time_dir = os.path.join(test.get_temp_dir(), 'time')
@@ -706,6 +709,7 @@
exception_str)
self.assertTrue(mat is None)
+ @test_util.run_v1_only('b/120545219')
def testTrackPersistentBytes(self):
ops.reset_default_graph()
a = array_ops.constant(np.ones((100, 100)))
diff --git a/tensorflow/python/profiler/pprof_profiler_test.py b/tensorflow/python/profiler/pprof_profiler_test.py
index 120a0d0..3f5bd9e 100644
--- a/tensorflow/python/profiler/pprof_profiler_test.py
+++ b/tensorflow/python/profiler/pprof_profiler_test.py
@@ -136,7 +136,7 @@
profile.ParseFromString(profile_contents)
self.assertEquals(expected_proto, str(profile))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testProfileWithWhileLoop(self):
options = config_pb2.RunOptions()
options.trace_level = config_pb2.RunOptions.FULL_TRACE
diff --git a/tensorflow/python/saved_model/BUILD b/tensorflow/python/saved_model/BUILD
index 53d0640..71d9e34 100644
--- a/tensorflow/python/saved_model/BUILD
+++ b/tensorflow/python/saved_model/BUILD
@@ -287,7 +287,7 @@
deps = [
":builder",
":constants",
- ":loader",
+ ":function_serialization",
":saved_object_graph_py",
":signature_constants",
":signature_def_utils",
@@ -295,15 +295,20 @@
":utils",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:array_ops",
+ "//tensorflow/python:constant_op",
+ "//tensorflow/python:control_flow_ops",
+ "//tensorflow/python:dtypes",
"//tensorflow/python:framework",
"//tensorflow/python:framework_ops",
"//tensorflow/python:lib",
"//tensorflow/python:resource_variable_ops",
+ "//tensorflow/python:tensor_spec",
"//tensorflow/python:util",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:function",
"//tensorflow/python/training/checkpointable:base",
+ "//tensorflow/python/training/checkpointable:tracking",
"//tensorflow/python/training/checkpointable:util",
],
)
@@ -330,8 +335,12 @@
],
srcs_version = "PY2AND3",
deps = [
+ ":constants",
+ ":function_deserialization",
":loader",
":saved_object_graph_py",
+ ":utils",
+ "//tensorflow/python:function",
"//tensorflow/python:lib",
"//tensorflow/python:util",
"//tensorflow/python/training/checkpointable:tracking",
@@ -345,10 +354,34 @@
deps = [
":load",
":save",
+ "//tensorflow/python:constant_op",
"//tensorflow/python:dtypes",
+ "//tensorflow/python:lib",
"//tensorflow/python:tensor_spec",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:test",
"//tensorflow/python/training/checkpointable:tracking",
],
)
+
+py_library(
+ name = "function_serialization",
+ srcs = [
+ "function_serialization.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":saved_object_graph_py",
+ "//tensorflow/python/eager:def_function",
+ "//tensorflow/python/eager:function",
+ ],
+)
+
+py_library(
+ name = "function_deserialization",
+ srcs = [
+ "function_deserialization.py",
+ ],
+ srcs_version = "PY2AND3",
+ deps = ["//tensorflow/python/eager:def_function"],
+)
diff --git a/tensorflow/python/saved_model/function_deserialization.py b/tensorflow/python/saved_model/function_deserialization.py
new file mode 100644
index 0000000..46bd69a
--- /dev/null
+++ b/tensorflow/python/saved_model/function_deserialization.py
@@ -0,0 +1,46 @@
+# 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.
+# ==============================================================================
+"""Tools for deserializing PolymorphicFunctions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.eager import def_function
+
+
+def recreate_polymorphic_function(
+ saved_polymorphic_function, defined_functions):
+ """Creates a PolymorphicFunction which runs restored function definitions."""
+ @def_function.function
+ def restored_function(*args):
+ """Calls a restored function."""
+ # Try calling each function, return a value from the first one whose
+ # signature matches.
+ # TODO(allenl): Consider re-populating the function cache directly.
+ # TODO(allenl): Functions saved with input_signatures should revive with
+ # input_signatures.
+ for monomorphic_function in saved_polymorphic_function.monomorphic_function:
+ try:
+ # TODO(allenl): Passing an explicit name here prevents invalid name
+ # errors. We should replace this with something based on the actual
+ # Python function name.
+ return defined_functions[monomorphic_function.concrete_function](
+ *args, name="imported_function")
+ except ValueError:
+ continue
+ raise AssertionError(
+ "Could not find matching function to call for arguments: %s" % (args,))
+ return restored_function
diff --git a/tensorflow/python/saved_model/function_serialization.py b/tensorflow/python/saved_model/function_serialization.py
new file mode 100644
index 0000000..69f34f0
--- /dev/null
+++ b/tensorflow/python/saved_model/function_serialization.py
@@ -0,0 +1,78 @@
+# 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.
+# ==============================================================================
+"""Tools for serializing PolymorphicFunctions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as defun_lib
+from tensorflow.python.saved_model import saved_object_graph_pb2
+
+
+def _serialize_polymorphic_function(polymorphic_function):
+ monomorphic_functions = []
+ for concrete_function in list_all_concrete_functions(polymorphic_function):
+ monomorphic_functions.append(
+ saved_object_graph_pb2.SavedMonomorphicFunction(
+ concrete_function=concrete_function.name))
+ saved_polymorphic_function = saved_object_graph_pb2.SavedPolymorphicFunction(
+ monomorphic_function=monomorphic_functions)
+ return saved_polymorphic_function
+
+
+def list_all_concrete_functions(polymorphic_function):
+ """Given a polymorphic function, returns all of its concrete functions."""
+ concrete_functions = []
+ for signature in polymorphic_function._cached_input_signatures: # pylint: disable=protected-access
+ if any(isinstance(arg, defun_lib.UnknownArgument) for arg in signature):
+ continue
+ concrete_function = polymorphic_function.get_concrete_function(*signature)
+ concrete_functions.append(concrete_function)
+ return concrete_functions
+
+
+def list_all_polymorphic_functions(checkpointable_object):
+ """Given a checkpointable object, returns all of its polymorphic functions."""
+ polymorphic_functions = dict()
+ for attribute_name in dir(checkpointable_object):
+ try:
+ attribute_value = getattr(checkpointable_object, attribute_name, None)
+ except: # pylint: disable=bare-except
+ # We really don't want to throw an exception just because some object's
+ # attribute accessor is broken.
+ attribute_value = None
+ # TODO(allenl): Consider de-duplicating functions which are referenced
+ # from multiple attributes.
+ if isinstance(attribute_value, def_function.PolymorphicFunction):
+ polymorphic_functions[attribute_name] = attribute_value
+ return polymorphic_functions
+
+
+def add_polymorphic_functions_to_object_graph_proto(checkpointable_objects,
+ saved_object_graph):
+ """Finds PolymorphicFunctions attached to objects and saves them."""
+ existing_objects = list(zip(checkpointable_objects, saved_object_graph.nodes))
+ for obj, obj_proto in existing_objects:
+ for name, polymorphic_function in list_all_polymorphic_functions(
+ obj).items():
+ function_node_id = len(saved_object_graph.nodes)
+ function_node = saved_object_graph.nodes.add()
+ function_node.function.CopyFrom(
+ _serialize_polymorphic_function(polymorphic_function))
+ reference = obj_proto.children.add()
+ reference.node_id = function_node_id
+ reference.local_name = name
diff --git a/tensorflow/python/saved_model/load.py b/tensorflow/python/saved_model/load.py
index 397086e..28c0af2 100644
--- a/tensorflow/python/saved_model/load.py
+++ b/tensorflow/python/saved_model/load.py
@@ -20,42 +20,91 @@
import os
+from tensorflow.python.framework import function as function_lib
from tensorflow.python.lib.io import file_io
from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import function_deserialization
+from tensorflow.python.saved_model import loader_impl
from tensorflow.python.saved_model import saved_object_graph_pb2
+from tensorflow.python.saved_model import utils_impl as saved_model_utils
from tensorflow.python.training.checkpointable import tracking
from tensorflow.python.util import compat
-def _recreate_object_graph(object_graph_proto):
- """Recreates Python objects from an ObjectGraph proto."""
- objects = []
- for _ in object_graph_proto.nodes:
- # TODO(allenl): re-create variables and other types
- objects.append(tracking.Checkpointable())
- for obj, object_proto in zip(objects, object_graph_proto.nodes):
- for reference in object_proto.children:
- setattr(obj, reference.local_name, objects[reference.node_id])
- return objects[0]
+class _Loader(object):
+ """Helper class to load an object-based SavedModel."""
+
+ def __init__(self, object_graph_proto, saved_model_proto, export_dir):
+ meta_graph = saved_model_proto.meta_graphs[0]
+ self._asset_file_def = meta_graph.asset_file_def
+ self._proto = object_graph_proto
+ self._export_dir = export_dir
+ self._defined_functions = {}
+ for defined_function in function_lib.from_library(
+ meta_graph.graph_def.library):
+ # TODO(allenl): Do we need to do name mapping here? Not quite sure what
+ # happens when loaded names collide with existing names.
+ defined_function.add_to_graph(None)
+ self._defined_functions[defined_function.name] = defined_function
+ self._load_all()
+
+ def _load_all(self):
+ self._nodes = [self._recreate(proto) for proto in self._proto.nodes]
+ # After creating the objects, construct the edges between the objects.
+ for obj, object_proto in zip(self._nodes, self._proto.nodes):
+ for reference in object_proto.children:
+ setattr(obj, reference.local_name, self._nodes[reference.node_id])
+
+ def get(self, node_id):
+ return self._nodes[node_id]
+
+ def _recreate(self, proto):
+ factory = {
+ "user_object": lambda: self._recreate_user_object(proto.user_object),
+ "asset": lambda: self._recreate_asset(proto.asset),
+ "function": lambda: self._recreate_function(proto.function)
+ }
+ kind = proto.WhichOneof("kind")
+ if kind not in factory:
+ raise ValueError("Unknown SavedObject type: %r" % kind)
+ return factory[kind]()
+
+ def _recreate_user_object(self, proto):
+ del proto
+ return tracking.Checkpointable()
+
+ def _recreate_asset(self, proto):
+ filename = os.path.join(
+ saved_model_utils.get_assets_dir(self._export_dir),
+ self._asset_file_def[proto.asset_file_def_index].filename)
+ return tracking.TrackableAsset(filename)
+
+ def _recreate_function(self, proto):
+ return function_deserialization.recreate_polymorphic_function(
+ proto, self._defined_functions)
+
+
+def _load_saved_object_graph_proto(filename):
+ with file_io.FileIO(filename, "rb") as f:
+ contents = f.read()
+ return saved_object_graph_pb2.SavedObjectGraph.FromString(contents)
def load(export_dir):
"""Load a SavedModel from `export_dir`."""
+ saved_model_proto = loader_impl.parse_saved_model(export_dir)
object_graph_filename = os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes(constants.EXTRA_ASSETS_DIRECTORY),
compat.as_bytes("object_graph.pb"))
if file_io.file_exists(object_graph_filename):
- # If there is an object graph associated with the SavedModel, we'll create a
- # root object from that.
- object_graph_string = file_io.FileIO(object_graph_filename, "rb").read()
- object_graph_proto = (
- saved_object_graph_pb2.SavedObjectGraph())
- object_graph_proto.ParseFromString(object_graph_string)
- root = _recreate_object_graph(object_graph_proto)
+ object_graph_proto = _load_saved_object_graph_proto(object_graph_filename)
+ loader = _Loader(object_graph_proto,
+ saved_model_proto,
+ export_dir)
+ root = loader.get(0)
else:
raise NotImplementedError(
"Currently only SavedModels exported with `tf.saved_model.save` may be "
"imported. Other SavedModels may eventually be supported via load().")
- # TODO(allenl): load functions from the SavedModel into the eager context
return root
diff --git a/tensorflow/python/saved_model/load_test.py b/tensorflow/python/saved_model/load_test.py
index bfa201a..303b8f6 100644
--- a/tensorflow/python/saved_model/load_test.py
+++ b/tensorflow/python/saved_model/load_test.py
@@ -19,11 +19,14 @@
from __future__ import print_function
import os
+import tempfile
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 tensor_spec
+from tensorflow.python.lib.io import file_io
from tensorflow.python.saved_model import load
from tensorflow.python.saved_model import save
from tensorflow.python.training.checkpointable import tracking
@@ -45,7 +48,73 @@
imported = load.load(save_dir)
self.assertIs(imported.dep_three, imported.dep_two.dep)
self.assertIsNot(imported.dep_one, imported.dep_two)
+ self.assertEqual(4., imported.f(constant_op.constant(2.)).numpy())
+ def _make_asset(self, contents):
+ filename = tempfile.mktemp(prefix=self.get_temp_dir())
+ with open(filename, "w") as f:
+ f.write(contents)
+ return filename
+
+ def test_assets_import(self):
+ file1 = self._make_asset("contents 1")
+ file2 = self._make_asset("contents 2")
+
+ root = tracking.Checkpointable()
+ root.f = def_function.function(
+ lambda x: 2. * x,
+ input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+ root.asset1 = tracking.TrackableAsset(file1)
+ root.asset2 = tracking.TrackableAsset(file2)
+
+ save_dir = os.path.join(self.get_temp_dir(), "save_dir")
+ save.save(root, save_dir)
+
+ file_io.delete_file(file1)
+ file_io.delete_file(file2)
+ load_dir = os.path.join(self.get_temp_dir(), "load_dir")
+ file_io.rename(save_dir, load_dir)
+
+ imported = load.load(load_dir)
+ with open(imported.asset1.asset_path.numpy(), "r") as f:
+ self.assertEquals("contents 1", f.read())
+ with open(imported.asset2.asset_path.numpy(), "r") as f:
+ self.assertEquals("contents 2", f.read())
+
+ def test_assets_dedup(self):
+ vocab = self._make_asset("contents")
+ root = tracking.Checkpointable()
+ root.f = def_function.function(
+ lambda x: 2. * x,
+ input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+
+ root.asset1 = tracking.TrackableAsset(vocab)
+ root.asset2 = tracking.TrackableAsset(vocab)
+
+ export_dir = os.path.join(self.get_temp_dir(), "save_dir")
+ save.save(root, export_dir)
+ imported = load.load(export_dir)
+
+ self.assertEqual(imported.asset1.asset_path.numpy(),
+ imported.asset2.asset_path.numpy())
+
+ def test_only_implicit_signatures(self):
+ def func(x):
+ return 2 * x
+
+ root = tracking.Checkpointable()
+ root.f = def_function.function(func)
+
+ # Add two traces.
+ root.f(constant_op.constant(1.))
+ root.f(constant_op.constant(1))
+
+ save_dir = os.path.join(self.get_temp_dir(), "saved_model")
+ save.save(root, save_dir, signatures=dict())
+ imported = load.load(save_dir)
+
+ self.assertEqual(4., imported.f(constant_op.constant(2.)).numpy())
+ self.assertEqual(14, imported.f(constant_op.constant(7)).numpy())
if __name__ == "__main__":
test.main()
diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py
index 6bf39a2..e5be03a 100644
--- a/tensorflow/python/saved_model/loader_impl.py
+++ b/tensorflow/python/saved_model/loader_impl.py
@@ -39,7 +39,7 @@
from tensorflow.python.util.tf_export import tf_export
-def _parse_saved_model(export_dir):
+def parse_saved_model(export_dir):
"""Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`.
Args:
@@ -83,6 +83,11 @@
constants.SAVED_MODEL_FILENAME_PB))
+# TODO(b/120594573): Make this symbol also available as private, so that
+# tensorflow_transform and tensorflow_estimator do not break.
+_parse_saved_model = parse_saved_model
+
+
def _get_asset_tensors(export_dir, meta_graph_def_to_load, import_scope=None):
"""Gets the asset tensors, if defined in the meta graph def to load.
@@ -276,7 +281,7 @@
"""
self._export_dir = export_dir
self._variables_path = saved_model_utils.get_variables_path(export_dir)
- self._saved_model = _parse_saved_model(export_dir)
+ self._saved_model = parse_saved_model(export_dir)
@property
def export_dir(self):
diff --git a/tensorflow/python/saved_model/loader_test.py b/tensorflow/python/saved_model/loader_test.py
index 3b7f0b2..3e27c08 100644
--- a/tensorflow/python/saved_model/loader_test.py
+++ b/tensorflow/python/saved_model/loader_test.py
@@ -94,7 +94,7 @@
super(SavedModelLoaderTest, self).tearDown()
shutil.rmtree(test.get_temp_dir(), ignore_errors=True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def test_load_function(self, builder_cls):
self.export_simple_graph(builder_cls)
loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
@@ -110,7 +110,7 @@
self.assertEqual(5, sess.graph.get_tensor_by_name("x:0").eval())
self.assertEqual(7, sess.graph.get_tensor_by_name("y:0").eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def test_load_graph(self, builder_cls):
self.export_simple_graph(builder_cls)
loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
@@ -130,7 +130,7 @@
with self.assertRaises(errors.FailedPreconditionError):
self.evaluate(y)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def test_load_with_import_scope(self, builder_cls):
self.export_graph_with_main_op(builder_cls)
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
@@ -179,7 +179,7 @@
loader.restore_variables(sess, tf_saver.Saver())
self.assertEqual(55, self.evaluate(z))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def test_run_init_op(self, builder_cls):
self.export_graph_with_main_op(builder_cls)
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
@@ -212,7 +212,7 @@
with self.assertRaises(RuntimeError):
loader.get_meta_graph_def_from_tags(["not_a_graph"])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def test_load_saved_model_with_no_variables(self, builder_cls):
"""Test that SavedModel runs saver when there appear to be no variables.
diff --git a/tensorflow/python/saved_model/save.py b/tensorflow/python/saved_model/save.py
index a66e19b..57c63f8 100644
--- a/tensorflow/python/saved_model/save.py
+++ b/tensorflow/python/saved_model/save.py
@@ -31,13 +31,13 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_spec
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.saved_model import builder_impl
from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import function_serialization
from tensorflow.python.saved_model import saved_object_graph_pb2
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils
@@ -50,28 +50,7 @@
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
-
-def _check_for_functional_keras_model(root):
- """Makes an export signature for `root` if it's a functional Keras Model."""
- # If nothing is decorated yet but this is a functional Keras Model (duck
- # typed), we'll try to make a signature ourselves.
- try:
- inputs = root.inputs
- input_names = root.input_names
- except AttributeError:
- return None
- input_signature = []
- for input_tensor, input_name in zip(inputs, input_names):
- input_signature.append(tensor_spec.TensorSpec(
- shape=input_tensor.shape, dtype=input_tensor.dtype,
- name=input_name))
-
- @def_function.function(input_signature=input_signature)
- def _wrapped_model(*args):
- outputs_list = nest.flatten(root(inputs=list(args)))
- return {name: output for name, output
- in zip(root.output_names, outputs_list)}
- return _wrapped_model
+DEFAULT_SIGNATURE_ATTR = "_default_save_signature"
def _find_function_to_export(root):
@@ -93,7 +72,7 @@
exported_function = attribute_value
previous_attribute_name = attribute_name
if exported_function is None:
- exported_function = _check_for_functional_keras_model(root)
+ exported_function = getattr(root, DEFAULT_SIGNATURE_ATTR, None)
if exported_function is None:
raise ValueError(
("Exporting an object with no tf.saved_model.save(..., signatures=...) "
@@ -375,7 +354,9 @@
# Map from asset variable resource Tensors to their init ops
"asset_initializers_by_resource",
# Map from base asset filenames to full paths
- "asset_filename_map"])
+ "asset_filename_map",
+ # Map from TrackableAsset to index of corresponding AssetFileDef
+ "asset_index"])
def _process_asset(trackable_asset, asset_info, resource_map):
@@ -386,21 +367,23 @@
path = builder_impl.get_asset_filename_to_add(
asset_filepath=original_path,
asset_filename_map=asset_info.asset_filename_map)
- asset_variable = asset_info.asset_filename_map.get(path, None)
- if asset_variable is None:
- asset_path_initializer = array_ops.placeholder(
- shape=original_variable.shape,
- dtype=dtypes.string,
- name="asset_path_initializer")
- asset_variable = resource_variable_ops.ResourceVariable(
- asset_path_initializer)
- asset_info.asset_filename_map[path] = original_path
- asset_def = meta_graph_pb2.AssetFileDef()
- asset_def.filename = path
- asset_def.tensor_info.name = asset_path_initializer.name
- asset_info.asset_defs.append(asset_def)
- asset_info.asset_initializers_by_resource[original_variable.handle] = (
- asset_variable.initializer)
+ # TODO(andresp): Instead of mapping 1-1 between trackable asset
+ # and asset in the graph def consider deduping the assets that
+ # point to the same file.
+ asset_path_initializer = array_ops.placeholder(
+ shape=original_variable.shape,
+ dtype=dtypes.string,
+ name="asset_path_initializer")
+ asset_variable = resource_variable_ops.ResourceVariable(
+ asset_path_initializer)
+ asset_info.asset_filename_map[path] = original_path
+ asset_def = meta_graph_pb2.AssetFileDef()
+ asset_def.filename = path
+ asset_def.tensor_info.name = asset_path_initializer.name
+ asset_info.asset_defs.append(asset_def)
+ asset_info.asset_initializers_by_resource[original_variable.handle] = (
+ asset_variable.initializer)
+ asset_info.asset_index[trackable_asset] = len(asset_info.asset_defs) - 1
resource_map[original_variable.handle] = asset_variable.handle
@@ -432,7 +415,8 @@
asset_info = _AssetInfo(
asset_defs=[],
asset_initializers_by_resource={},
- asset_filename_map={})
+ asset_filename_map={},
+ asset_index={})
for obj in accessible_objects:
if isinstance(obj, tracking.TrackableResource):
new_resource = obj.create_resource()
@@ -441,7 +425,7 @@
new_variable = resource_variable_ops.copy_to_graph_uninitialized(obj)
object_map[obj] = new_variable
resource_map[obj.handle] = new_variable.handle
- if isinstance(obj, tracking.TrackableAsset):
+ elif isinstance(obj, tracking.TrackableAsset):
_process_asset(obj, asset_info, resource_map)
return object_map, resource_map, asset_info
@@ -458,9 +442,7 @@
object_saver: A CheckpointableSaver to add to the MetaGraph.
Returns:
- asset_filename_map, a dictionary mapping from asset base names to
- user-specified full asset paths, which should be copied to the SavedModel's
- assets/ directory.
+ An _AssetInfo, which contains information to help creating the SavedModel.
"""
signatures = {}
# List objects from the eager context to make sure Optimizers give us the
@@ -499,8 +481,20 @@
# variables, but want any operations associated with the save/restore to be in
# the exported graph (thus the `to_graph` argument).
saver = object_saver.freeze(object_map=object_map, to_graph=exported_graph)
+
+ # We must resolve the concrete function to add to MetaGraph while in eager
+ # mode.
+ concrete_functions = []
+ for accessible_object in accessible_objects:
+ for function in function_serialization.list_all_polymorphic_functions(
+ accessible_object).values():
+ concrete_functions.extend(
+ function_serialization.list_all_concrete_functions(function))
+
with exported_graph.as_default():
signatures = _generate_signatures(signature_functions, resource_map)
+ for concrete_function in concrete_functions:
+ concrete_function.add_to_graph()
saver_def = saver.to_proto()
meta_graph_def.saver_def.CopyFrom(saver_def)
graph_def = exported_graph.as_graph_def(add_shapes=True)
@@ -514,23 +508,41 @@
for signature_key, signature in signatures.items():
meta_graph_def.signature_def[signature_key].CopyFrom(signature)
meta_graph.strip_graph_default_valued_attrs(meta_graph_def)
- return asset_info.asset_filename_map
+ return asset_info
-def _write_object_graph(obj, export_dir):
- """Save a SavedObjectGraph proto for `obj`."""
+def _write_object_graph(root, export_dir, asset_file_def_index):
+ """Save a SavedObjectGraph proto for `root`."""
# SavedObjectGraph is similar to the CheckpointableObjectGraph proto in the
# checkpoint. It will eventually go into the SavedModel.
- object_proto = util.make_object_graph_without_attributes(
- obj, proto=saved_object_graph_pb2.SavedObjectGraph())
+ proto = saved_object_graph_pb2.SavedObjectGraph()
+
+ checkpointable_objects, node_ids, slot_variables = util.find_objects(root)
+ util.fill_object_graph_proto(checkpointable_objects, node_ids, slot_variables,
+ proto)
+
+ for obj, obj_proto in zip(checkpointable_objects, proto.nodes):
+ _write_object_proto(obj, obj_proto, asset_file_def_index)
+
+ function_serialization.add_polymorphic_functions_to_object_graph_proto(
+ checkpointable_objects, proto)
+
extra_asset_dir = os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes(constants.EXTRA_ASSETS_DIRECTORY))
file_io.recursive_create_dir(extra_asset_dir)
object_graph_filename = os.path.join(
extra_asset_dir, compat.as_bytes("object_graph.pb"))
- file_io.write_string_to_file(object_graph_filename,
- object_proto.SerializeToString())
+ file_io.write_string_to_file(object_graph_filename, proto.SerializeToString())
+
+
+def _write_object_proto(obj, proto, asset_file_def_index):
+ """Saves an object into SavedObject proto."""
+ if isinstance(obj, tracking.TrackableAsset):
+ proto.asset.SetInParent()
+ proto.asset.asset_file_def_index = asset_file_def_index[obj]
+ else:
+ proto.user_object.SetInParent()
@tf_export("saved_model.save", v1=["saved_model.experimental.save"])
@@ -680,7 +692,25 @@
Raises:
ValueError: If `obj` is not checkpointable.
+
+ @compatibility(eager)
+ Not supported when graph building. From TensorFlow 1.x,
+ `tf.enable_eager_execution()` must run first. May not be called from within a
+ function body.
+ @end_compatibility
"""
+ if not context.executing_eagerly():
+ with ops.init_scope():
+ if context.executing_eagerly():
+ raise AssertionError(
+ "tf.saved_model.save is not supported inside a traced "
+ "@tf.function. Move the call to the outer eagerly-executed "
+ "context.")
+ else:
+ raise AssertionError(
+ "tf.saved_model.save is not supported when graph building. "
+ "tf.enable_eager_execution() must run first when calling it from "
+ "TensorFlow 1.x.")
# pylint: enable=line-too-long
if not isinstance(obj, base.CheckpointableBase):
raise ValueError(
@@ -697,7 +727,7 @@
saved_model = saved_model_pb2.SavedModel()
meta_graph_def = saved_model.meta_graphs.add()
object_saver = util.CheckpointableSaver(obj)
- asset_filename_map = _fill_meta_graph_def(
+ asset_info = _fill_meta_graph_def(
meta_graph_def, obj, signatures, object_saver)
saved_model.saved_model_schema_version = (
constants.SAVED_MODEL_SCHEMA_VERSION)
@@ -706,9 +736,10 @@
# SavedModel proto itself.
utils_impl.get_or_create_variables_dir(export_dir)
object_saver.save(utils_impl.get_variables_path(export_dir))
- builder_impl.copy_assets_to_destination_dir(asset_filename_map, export_dir)
+ builder_impl.copy_assets_to_destination_dir(asset_info.asset_filename_map,
+ export_dir)
path = os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
file_io.write_string_to_file(path, saved_model.SerializeToString())
- _write_object_graph(obj, export_dir)
+ _write_object_graph(obj, export_dir, asset_info.asset_index)
diff --git a/tensorflow/python/saved_model/save_test.py b/tensorflow/python/saved_model/save_test.py
index 5096e9c..5381c2f 100644
--- a/tensorflow/python/saved_model/save_test.py
+++ b/tensorflow/python/saved_model/save_test.py
@@ -21,8 +21,6 @@
import os
import sys
-import numpy
-
from tensorflow.python.client import session as session_lib
from tensorflow.python.eager import backprop
from tensorflow.python.eager import def_function
@@ -32,12 +30,8 @@
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
-from tensorflow.python.keras.engine import input_layer
-from tensorflow.python.keras.engine import training
from tensorflow.python.keras.layers import core
-from tensorflow.python.keras.layers import merge
from tensorflow.python.lib.io import file_io
-from tensorflow.python.ops import array_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
@@ -50,10 +44,9 @@
from tensorflow.python.training.checkpointable import util
-class _ModelWithOptimizer(training.Model):
+class _ModelWithOptimizer(util.Checkpoint):
def __init__(self):
- super(_ModelWithOptimizer, self).__init__()
self.dense = core.Dense(1)
self.optimizer = adam.AdamOptimizer(0.01)
@@ -63,7 +56,7 @@
def call(self, x, y):
with backprop.GradientTape() as tape:
loss = math_ops.reduce_mean((self.dense(x) - y) ** 2.)
- trainable_variables = self.trainable_variables
+ trainable_variables = self.dense.trainable_variables
gradients = tape.gradient(loss, trainable_variables)
self.optimizer.apply_gradients(zip(gradients, trainable_variables))
return {"loss": loss}
@@ -179,10 +172,10 @@
x = constant_op.constant([[3., 4.]])
y = constant_op.constant([2.])
model = _ModelWithOptimizer()
- first_loss = model(x, y)
+ first_loss = model.call(x, y)
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
save.save(model, save_dir, model.call)
- second_loss = model(x, y)
+ second_loss = model.call(x, y)
self.assertNotEqual(first_loss, second_loss)
self.assertAllClose(
second_loss,
@@ -197,7 +190,7 @@
model = _ModelWithOptimizer()
x = constant_op.constant([[3., 4.]])
y = constant_op.constant([2.])
- model(x, y)
+ model.call(x, y)
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
save.save(model, save_dir)
self.assertIn("loss",
@@ -217,25 +210,40 @@
model = _ModelWithOptimizer()
x = constant_op.constant([[3., 4.]])
y = constant_op.constant([2.])
- model(x, y)
+ model.call(x, y)
model.second_function = def_function.function(lambda: 1.)
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
with self.assertRaisesRegexp(ValueError, "call.*second_function"):
save.save(model, save_dir)
- def test_subclassed_no_signature(self):
+ def test_no_signature(self):
- class Subclassed(training.Model):
+ class Model(util.Checkpoint):
def call(self, inputs):
return inputs * 2.
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
- model = Subclassed()
+ model = Model()
with self.assertRaisesRegexp(
ValueError, "no @tf.function-decorated methods"):
save.save(model, save_dir)
+ def test_find_default_save_function(self):
+
+ class ObjWithDefaultSignature(util.Checkpoint):
+
+ @def_function.function(input_signature=[tensor_spec.TensorSpec(
+ shape=None, dtype=dtypes.float32)])
+ def _default_save_signature(self, x):
+ return x + x + 1
+
+ obj = ObjWithDefaultSignature()
+ save_dir = os.path.join(self.get_temp_dir(), "saved_model")
+ save.save(obj, save_dir)
+ self.assertAllClose(
+ {"output_0": 7.}, _import_and_infer(save_dir, {"x": 3.}))
+
def test_docstring(self):
class Adder(util.Checkpoint):
@@ -276,46 +284,6 @@
self.assertNotIn("T", complex_node.attr)
self.assertNotIn("Tout", complex_node.attr)
- def test_export_functional_keras_model(self):
- x = input_layer.Input((4,), name="x")
- y = core.Dense(4, name="out")(x)
- model = training.Model(x, y)
- save_dir = os.path.join(self.get_temp_dir(), "saved_model")
- save.save(model, save_dir)
- self.assertAllClose(
- {"out": model(array_ops.ones([1, 4]))},
- _import_and_infer(save_dir, {"x": [[1., 1., 1., 1.]]}))
-
- @test_util.run_deprecated_v1
- def test_export_functional_keras_model_after_fit(self):
- x = input_layer.Input((1,))
- y = core.Dense(1, name="y")(x)
- model = training.Model(x, y)
- model.compile(optimizer="sgd", loss="mse")
- model.fit(x=numpy.array([[1.]]),
- y=numpy.array([2.]), epochs=2)
- save_dir = os.path.join(self.get_temp_dir(), "saved_model")
- save.save(model, save_dir)
- self.assertAllClose(
- {"y": model(constant_op.constant([[1.], [2.]]))},
- _import_and_infer(save_dir, {"input_1": [[1.], [2.]]}))
-
- def test_export_multi_input_functional_keras_model(self):
- x1 = input_layer.Input((2,), name="x1")
- x2 = input_layer.Input((2,), name="x2")
- y1 = core.Dense(4)(merge.Add()([x1, x2]))
- y2 = core.Dense(4)(merge.Multiply()([x1, x2]))
- model = training.Model([x1, x2], [y1, y2])
- save_dir = os.path.join(self.get_temp_dir(), "saved_model")
- save.save(model, save_dir)
- outputs = model([array_ops.ones([1, 2]), 2. * array_ops.ones([1, 2])])
- self.assertAllClose(
- {"dense": outputs[0], "dense_1": outputs[1]},
- _import_and_infer(
- save_dir,
- {"x1": [[1., 1.]],
- "x2": [[2., 2.]]}))
-
class AssetTests(test.TestCase):
@@ -353,6 +321,34 @@
{"output_0": [2, 1]},
_import_and_infer(second_dir, {"keys": ["gamma", "beta"]}))
+ def test_unused_asset(self):
+ root = tracking.Checkpointable()
+ root.f = def_function.function(
+ lambda x: 2. * x,
+ input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+ root.asset = tracking.TrackableAsset(self._vocab_path)
+
+ export_dir = os.path.join(self.get_temp_dir(), "save_dir")
+ save.save(root, export_dir)
+ self.assertAllClose(
+ {"output_0": [0.2]},
+ _import_and_infer(export_dir, {"x": [0.1]}))
+
+ def test_sensible_graph_building_exception(self):
+ root = util.Checkpoint(v=variables.Variable(2.))
+ root.f = def_function.function(
+ lambda x: 2. * root.v,
+ input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+ export_dir = os.path.join(self.get_temp_dir(), "save_dir")
+ @def_function.function
+ def _calls_save():
+ save.save(root, export_dir)
+ with self.assertRaisesRegexp(AssertionError, "tf.function"):
+ _calls_save()
+ with ops.Graph().as_default():
+ with self.assertRaisesRegexp(AssertionError, "enable_eager_execution"):
+ save.save(root, export_dir)
+
class MemoryTests(test.TestCase):
@@ -363,7 +359,7 @@
def test_no_reference_cycles(self):
x = constant_op.constant([[3., 4.]])
y = constant_op.constant([2.])
- self._model(x, y)
+ self._model.call(x, y)
if sys.version_info[0] < 3:
# TODO(allenl): debug reference cycles in Python 2.x
self.skipTest("This test only works in Python 3+. Reference cycles are "
diff --git a/tensorflow/python/saved_model/saved_model_test.py b/tensorflow/python/saved_model/saved_model_test.py
index 0f18fb1..e36b8b3 100644
--- a/tensorflow/python/saved_model/saved_model_test.py
+++ b/tensorflow/python/saved_model/saved_model_test.py
@@ -328,7 +328,7 @@
self.assertRaises(RuntimeError, loader.load, sess, ["foo", "baz"],
export_dir)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testVariables(self):
export_dir = self._get_export_dir("test_variables")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -474,7 +474,7 @@
self.assertEqual(
42, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCollections(self):
export_dir = self._get_export_dir("test_collections")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -819,6 +819,7 @@
self._validate_assets(export_dir, foo_graph.asset_file_def, "hello42.txt",
"foo bar baz 0", "asset_file_tensor_0:0")
+ @test_util.run_v1_only("b/120545219")
def testCustomInitOp(self):
export_dir = self._get_export_dir("test_main_op")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -854,7 +855,7 @@
# the main_op, following a restore.
self.assertEqual(3, ops.get_collection("v")[2].eval())
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testTrainOp(self):
export_dir = self._get_export_dir("test_train_op")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -882,7 +883,7 @@
self.assertIsInstance(
loader_impl.get_train_op(meta_graph_def), ops.Tensor)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testTrainOpGroup(self):
export_dir = self._get_export_dir("test_train_op_group")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -910,7 +911,7 @@
self.assertIsInstance(
loader_impl.get_train_op(meta_graph_def), ops.Operation)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testTrainOpAfterVariables(self):
export_dir = self._get_export_dir("test_train_op_after_variables")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -1029,7 +1030,7 @@
self._validate_assets(export_dir, bar_graph.asset_file_def, "foo.txt",
"content_foo", "asset_file_tensor:0")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testOp(self):
export_dir = self._get_export_dir("test_op")
builder = saved_model_builder._SavedModelBuilder(export_dir)
@@ -1083,7 +1084,7 @@
# CheckpointedOp is a key-value table that can be saved across sessions.
# The table register itself in SAVEABLE_OBJECTS collection.
v1 = saver_test_utils.CheckpointedOp(name="v1")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
v1.insert("k1", 3.0).run()
# Once the table is restored, we can access it through this reference.
ops.add_to_collection("table_ref", v1.table_ref)
@@ -1460,10 +1461,8 @@
self.assertIn("Tout", complex_node.attr)
# Load graph "foo" from disk as-is to verify default attrs are stripped.
- # pylint: disable=protected-access
- saved_model_pb = loader_impl._parse_saved_model(export_dir)
+ saved_model_pb = loader_impl.parse_saved_model(export_dir)
self.assertIsNotNone(saved_model_pb)
- # pylint: enable=protected-access
meta_graph_foo_def = None
meta_graph_bar_def = None
@@ -1494,7 +1493,7 @@
self.assertIn("T", node_def.attr)
self.assertIn("Tout", node_def.attr)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testLegacyInitOp(self):
export_dir = self._get_export_dir("test_legacy_init_op")
builder = saved_model_builder.SavedModelBuilder(export_dir)
diff --git a/tensorflow/python/saved_model/saved_object_graph.proto b/tensorflow/python/saved_model/saved_object_graph.proto
index 89f82b9..ed5c639 100644
--- a/tensorflow/python/saved_model/saved_object_graph.proto
+++ b/tensorflow/python/saved_model/saved_object_graph.proto
@@ -18,21 +18,67 @@
// NOTE: This protocol buffer format is experimental and subject to change.
message SavedObjectGraph {
- message SavedObject {
- // Objects which this object depends on: named edges in the dependency
- // graph.
- repeated CheckpointableObjectGraph.CheckpointableObject.ObjectReference
- children = 1;
- // Removed when forking from CheckpointableObjectGraph.
- reserved "attributes";
- reserved 2;
- // Slot variables owned by this object. This describes the three-way
- // (optimizer, variable, slot variable) relationship; none of the three
- // depend on the others directly.
- repeated
- CheckpointableObjectGraph.CheckpointableObject.SlotVariableReference
- slot_variables = 3;
- }
-
+ // List of objects in the SavedModel.
+ //
+ // The position of the object in this list indicates its id.
+ // Nodes[0] is considered the root node.
repeated SavedObject nodes = 1;
}
+
+message SavedObject {
+ // Objects which this object depends on: named edges in the dependency
+ // graph.
+ //
+ // Note: only valid if kind == "object".
+ repeated CheckpointableObjectGraph.CheckpointableObject.ObjectReference
+ children = 1;
+
+ // Removed when forking from CheckpointableObjectGraph.
+ reserved "attributes";
+ reserved 2;
+
+ // Slot variables owned by this object. This describes the three-way
+ // (optimizer, variable, slot variable) relationship; none of the three
+ // depend on the others directly.
+ //
+ // Note: only valid if kind == "object".
+ repeated CheckpointableObjectGraph.CheckpointableObject.SlotVariableReference
+ slot_variables = 3;
+
+ oneof kind {
+ SavedUserObject user_object = 4;
+ SavedAsset asset = 5;
+ SavedPolymorphicFunction function = 6;
+ }
+}
+
+// A SavedUserObject is an object (in the object-oriented language of the
+// TensorFlow program) of some user- or framework-defined class other than
+// those handled specifically by the other kinds of SavedObjects.
+//
+// This object cannot be evaluated as a tensor, and therefore cannot be bound
+// to an input of a function.
+message SavedUserObject {}
+
+// A SavedAsset represents a file in a SavedModel.
+//
+// When bound to a function this object evaluates to a Variable from which the
+// absolute filename can be read. Users should not expect the filename to be
+// maintained.
+message SavedAsset {
+ // Index into `MetaGraphDef.asset_file_def[]` that describes the Asset.
+ //
+ // Only the field `AssetFileDef.filename` is used. Other fields, such as
+ // `AssetFileDef.tensor_info`, MUST be ignored.
+ uint32 asset_file_def_index = 1;
+}
+
+// A function with multiple signatures, possibly with non-Tensor arguments.
+message SavedPolymorphicFunction {
+ repeated SavedMonomorphicFunction monomorphic_function = 1;
+}
+
+message SavedMonomorphicFunction {
+ // A reference to a TensorFlow function in the MetaGraph's FunctionDefLibrary
+ string concrete_function = 1;
+}
diff --git a/tensorflow/python/saved_model/signature_def_utils_test.py b/tensorflow/python/saved_model/signature_def_utils_test.py
index 53c4523..d1347eb 100644
--- a/tensorflow/python/saved_model/signature_def_utils_test.py
+++ b/tensorflow/python/saved_model/signature_def_utils_test.py
@@ -423,6 +423,7 @@
{},
signature_constants.PREDICT_METHOD_NAME)
+ @test_util.run_v1_only("b/120545219")
def testOpSignatureDef(self):
key = "adding_1_and_2_key"
add_op = math_ops.add(1, 2, name="adding_1_and_2")
@@ -430,6 +431,7 @@
self.assertIn(key, signature_def.outputs)
self.assertEqual(add_op.name, signature_def.outputs[key].name)
+ @test_util.run_v1_only("b/120545219")
def testLoadOpFromSignatureDef(self):
key = "adding_1_and_2_key"
add_op = math_ops.add(1, 2, name="adding_1_and_2")
diff --git a/tensorflow/python/saved_model/utils_test.py b/tensorflow/python/saved_model/utils_test.py
index 0888dcb..2afe8ab 100644
--- a/tensorflow/python/saved_model/utils_test.py
+++ b/tensorflow/python/saved_model/utils_test.py
@@ -24,6 +24,7 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
+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.platform import test
@@ -32,6 +33,7 @@
class UtilsTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testBuildTensorInfoOp(self):
x = constant_op.constant(1, name="x")
y = constant_op.constant(2, name="y")
@@ -41,6 +43,7 @@
self.assertEqual(types_pb2.DT_INVALID, z_op_info.dtype)
self.assertEqual(0, len(z_op_info.tensor_shape.dim))
+ @test_util.run_v1_only("b/120545219")
def testBuildTensorInfoDefunOp(self):
@function.defun
def my_init_fn(x, y):
@@ -54,6 +57,7 @@
self.assertEqual(types_pb2.DT_INVALID, init_op_info.dtype)
self.assertEqual(0, len(init_op_info.tensor_shape.dim))
+ @test_util.run_v1_only("b/120545219")
def testBuildTensorInfoDense(self):
x = array_ops.placeholder(dtypes.float32, 1, name="x")
x_tensor_info = utils.build_tensor_info(x)
@@ -62,6 +66,7 @@
self.assertEqual(1, len(x_tensor_info.tensor_shape.dim))
self.assertEqual(1, x_tensor_info.tensor_shape.dim[0].size)
+ @test_util.run_v1_only("b/120545219")
def testBuildTensorInfoSparse(self):
x = array_ops.sparse_placeholder(dtypes.float32, [42, 69], name="x")
x_tensor_info = utils.build_tensor_info(x)
@@ -76,6 +81,7 @@
self.assertEqual(42, x_tensor_info.tensor_shape.dim[0].size)
self.assertEqual(69, x_tensor_info.tensor_shape.dim[1].size)
+ @test_util.run_v1_only("b/120545219")
def testGetTensorFromInfoDense(self):
expected = array_ops.placeholder(dtypes.float32, 1, name="x")
tensor_info = utils.build_tensor_info(expected)
@@ -83,6 +89,7 @@
self.assertIsInstance(actual, ops.Tensor)
self.assertEqual(expected.name, actual.name)
+ @test_util.run_v1_only("b/120545219")
def testGetTensorFromInfoSparse(self):
expected = array_ops.sparse_placeholder(dtypes.float32, name="x")
tensor_info = utils.build_tensor_info(expected)
@@ -122,6 +129,7 @@
import_scope="foo")
self.assertEqual(expected.name, actual.name)
+ @test_util.run_v1_only("b/120545219")
def testGetTensorFromInfoRaisesErrors(self):
expected = array_ops.placeholder(dtypes.float32, 1, name="x")
tensor_info = utils.build_tensor_info(expected)
diff --git a/tensorflow/python/tools/api/generator/api_init_files.bzl b/tensorflow/python/tools/api/generator/api_init_files.bzl
index 3517c11..5fee9c5 100644
--- a/tensorflow/python/tools/api/generator/api_init_files.bzl
+++ b/tensorflow/python/tools/api/generator/api_init_files.bzl
@@ -4,6 +4,8 @@
TENSORFLOW_API_INIT_FILES = [
# BEGIN GENERATED FILES
"__init__.py",
+ "autograph/__init__.py",
+ "autograph/experimental/__init__.py",
"bitwise/__init__.py",
"compat/__init__.py",
"data/__init__.py",
@@ -18,6 +20,7 @@
"graph_util/__init__.py",
"image/__init__.py",
"io/__init__.py",
+ "queue/__init__.py",
"initializers/__init__.py",
"keras/__init__.py",
"keras/activations/__init__.py",
@@ -64,10 +67,10 @@
"lite/constants/__init__.py",
"losses/__init__.py",
"math/__init__.py",
- "metrics/__init__.py",
"nn/__init__.py",
"nn/rnn_cell/__init__.py",
"quantization/__init__.py",
+ "ragged/__init__.py",
"random/__init__.py",
"saved_model/__init__.py",
"sets/__init__.py",
diff --git a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl
index e35b9c4..8d3b86b 100644
--- a/tensorflow/python/tools/api/generator/api_init_files_v1.bzl
+++ b/tensorflow/python/tools/api/generator/api_init_files_v1.bzl
@@ -5,6 +5,8 @@
# BEGIN GENERATED FILES
"__init__.py",
"app/__init__.py",
+ "autograph/__init__.py",
+ "autograph/experimental/__init__.py",
"bitwise/__init__.py",
"compat/__init__.py",
"data/__init__.py",
@@ -21,6 +23,7 @@
"graph_util/__init__.py",
"image/__init__.py",
"io/__init__.py",
+ "queue/__init__.py",
"initializers/__init__.py",
"keras/__init__.py",
"keras/activations/__init__.py",
@@ -77,6 +80,7 @@
"profiler/__init__.py",
"python_io/__init__.py",
"quantization/__init__.py",
+ "ragged/__init__.py",
"random/__init__.py",
"resource_loader/__init__.py",
"strings/__init__.py",
diff --git a/tensorflow/python/tools/api/generator/doc_srcs.py b/tensorflow/python/tools/api/generator/doc_srcs.py
index abb5886..b567eea 100644
--- a/tensorflow/python/tools/api/generator/doc_srcs.py
+++ b/tensorflow/python/tools/api/generator/doc_srcs.py
@@ -54,6 +54,7 @@
'nn': DocSource(docstring_module_name='ops.nn_ops'),
'nn.rnn_cell': DocSource(docstring_module_name='ops.rnn_cell'),
'python_io': DocSource(docstring_module_name='lib.io.python_io'),
+ 'ragged': DocSource(docstring_module_name='ops.ragged'),
'resource_loader': DocSource(
docstring_module_name='platform.resource_loader'),
'sets': DocSource(docstring_module_name='ops.sets'),
diff --git a/tensorflow/python/tools/freeze_graph_test.py b/tensorflow/python/tools/freeze_graph_test.py
index efdf7dd..de2672d 100644
--- a/tensorflow/python/tools/freeze_graph_test.py
+++ b/tensorflow/python/tools/freeze_graph_test.py
@@ -161,11 +161,11 @@
},)
builder.save(as_text=True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testFreezeGraphV1(self):
self._testFreezeGraph(saver_pb2.SaverDef.V1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testFreezeGraphV2(self):
self._testFreezeGraph(saver_pb2.SaverDef.V2)
diff --git a/tensorflow/python/tools/saved_model_cli.py b/tensorflow/python/tools/saved_model_cli.py
index c4c3756..afc4e51 100644
--- a/tensorflow/python/tools/saved_model_cli.py
+++ b/tensorflow/python/tools/saved_model_cli.py
@@ -659,6 +659,28 @@
scan_meta_graph_def(meta_graph_def)
+def convert_with_tensorrt(args):
+ """Function triggered by 'convert tensorrt' command.
+
+ Args:
+ args: A namespace parsed from command line.
+ """
+ # Import here instead of at top, because this will crash if TensorRT is
+ # not installed
+ from tensorflow.contrib import tensorrt # pylint: disable=g-import-not-at-top
+ tensorrt.create_inference_graph(
+ None,
+ None,
+ max_batch_size=args.max_batch_size,
+ max_workspace_size_bytes=args.max_workspace_size_bytes,
+ precision_mode=args.precision_mode,
+ minimum_segment_size=args.minimum_segment_size,
+ is_dynamic_op=args.is_dynamic_op,
+ input_saved_model_dir=args.dir,
+ input_saved_model_tags=args.tag_set.split(','),
+ output_saved_model_dir=args.output_dir)
+
+
def create_parser():
"""Creates a parser that parse the command line arguments.
@@ -812,6 +834,71 @@
help='tag-set of graph in SavedModel to scan, separated by \',\'')
parser_scan.set_defaults(func=scan)
+ # convert command
+ convert_msg = ('Usage example:\n'
+ 'To convert the SavedModel to one that have TensorRT ops:\n'
+ '$saved_model_cli convert \\\n'
+ ' --dir /tmp/saved_model \\\n'
+ ' --tag_set serve \\\n'
+ ' --output_dir /tmp/saved_model_trt \\\n'
+ ' tensorrt \n')
+ parser_convert = subparsers.add_parser(
+ 'convert',
+ description=convert_msg,
+ formatter_class=argparse.RawTextHelpFormatter)
+ parser_convert.add_argument(
+ '--dir',
+ type=str,
+ required=True,
+ help='directory containing the SavedModel to convert')
+ parser_convert.add_argument(
+ '--output_dir',
+ type=str,
+ required=True,
+ help='output directory for the converted SavedModel')
+ parser_convert.add_argument(
+ '--tag_set',
+ type=str,
+ required=True,
+ help='tag-set of graph in SavedModel to convert, separated by \',\'')
+ convert_subparsers = parser_convert.add_subparsers(
+ title='conversion methods',
+ description='valid conversion methods',
+ help='the conversion to run with the SavedModel')
+ parser_convert_with_tensorrt = convert_subparsers.add_parser(
+ 'tensorrt',
+ description='Convert the SavedModel with Tensorflow-TensorRT integration',
+ formatter_class=argparse.RawTextHelpFormatter)
+ parser_convert_with_tensorrt.add_argument(
+ '--max_batch_size',
+ type=int,
+ default=1,
+ help='max size for the input batch')
+ parser_convert_with_tensorrt.add_argument(
+ '--max_workspace_size_bytes',
+ type=int,
+ default=2 << 20,
+ help=('the maximum GPU temporary memory which the TRT engine can use at '
+ 'execution time'))
+ parser_convert_with_tensorrt.add_argument(
+ '--precision_mode',
+ type=str,
+ default='FP32',
+ help='one of FP32, FP16 and INT8')
+ parser_convert_with_tensorrt.add_argument(
+ '--minimum_segment_size',
+ type=int,
+ default=3,
+ help=('the minimum number of nodes required for a subgraph to be replaced'
+ 'in a TensorRT node'))
+ parser_convert_with_tensorrt.add_argument(
+ '--is_dynamic_op',
+ type=bool,
+ default=False,
+ help=('whether to generate dynamic TRT ops which will build the TRT '
+ 'network and engine at run time'))
+ parser_convert_with_tensorrt.set_defaults(func=convert_with_tensorrt)
+
return parser
diff --git a/tensorflow/python/training/adagrad_test.py b/tensorflow/python/training/adagrad_test.py
index da26fcd..3528fda 100644
--- a/tensorflow/python/training/adagrad_test.py
+++ b/tensorflow/python/training/adagrad_test.py
@@ -106,7 +106,7 @@
pred = math_ops.matmul(embedding_ops.embedding_lookup([var0], [0]), x)
loss = pred * pred
sgd_op = adagrad.AdagradOptimizer(1.0).minimize(loss)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Fetch params to validate initial values
self.assertAllCloseAccordingToType([[1.0, 2.0], [3.0, 4.0]],
self.evaluate(var0))
@@ -129,7 +129,7 @@
constant_op.constant(3.0), initial_accumulator_value=0.1)
ada_update = ada_opt.apply_gradients(
zip([grads0, grads1], [var0, var1]))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Fetch params to validate initial values
self.assertAllClose([1.0, 2.0], self.evaluate(var0))
self.assertAllClose([3.0, 4.0], self.evaluate(var1))
@@ -163,7 +163,7 @@
ada_opt = adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1)
ada_update = ada_opt.apply_gradients(
zip([grads0, grads1], [var0, var1]))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Fetch params to validate initial values
self.assertAllClose([[1.0], [2.0]], self.evaluate(var0))
self.assertAllClose([[3.0], [4.0]], self.evaluate(var1))
@@ -198,7 +198,7 @@
[(grad_repeated_index, repeated_index_update_var)])
aggregated_update = adagrad.AdagradOptimizer(3.0).apply_gradients(
[(grad_aggregated, aggregated_update_var)])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(aggregated_update_var.eval(),
self.evaluate(repeated_index_update_var))
for _ in range(3):
@@ -223,7 +223,7 @@
2.0).minimize(loss_repeated)
update_op_aggregated = adagrad.AdagradOptimizer(
2.0).minimize(loss_aggregated)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(
self.evaluate(var_repeated), self.evaluate(var_aggregated))
for _ in range(3):
@@ -289,7 +289,7 @@
self.assertEquals(slot0.get_shape(), var0.get_shape())
slot1 = ada_opt.get_slot(var1, "accumulator")
self.assertEquals(slot1.get_shape(), var1.get_shape())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Fetch params to validate initial values.
self.assertAllClose([1.0, 2.0], self.evaluate(var0))
@@ -306,7 +306,7 @@
np.array([2.715679168701172, 3.715679168701172]),
self.evaluate(var1))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testDynamicShapeVariable_Ok(self):
with self.cached_session():
v = variable_scope.get_variable("v", initializer=constant_op.constant(1.),
@@ -315,7 +315,7 @@
# Creating optimizer should cause no exception.
adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testDynamicShapeVariableWithCallableInit(self):
var0 = variable_scope.get_variable("var0",
initializer=constant_op.constant(1.),
diff --git a/tensorflow/python/training/adam_test.py b/tensorflow/python/training/adam_test.py
index b0bae27..1595811 100644
--- a/tensorflow/python/training/adam_test.py
+++ b/tensorflow/python/training/adam_test.py
@@ -68,8 +68,8 @@
var0 = resource_variable_ops.ResourceVariable(var0_np)
var1 = resource_variable_ops.ResourceVariable(var1_np)
else:
- var0 = variables.Variable(var0_np)
- var1 = variables.Variable(var1_np)
+ var0 = variables.RefVariable(var0_np)
+ var1 = variables.RefVariable(var1_np)
grads0_np_indices = np.array([0, 1], dtype=np.int32)
grads0 = ops.IndexedSlices(
constant_op.constant(grads0_np),
@@ -156,6 +156,9 @@
self.evaluate(repeated_index_update_var))
def doTestBasic(self, use_resource=False, use_callable_params=False):
+ if context.executing_eagerly() and not use_resource:
+ self.skipTest(
+ "Skipping test with use_resource=False and executing eagerly.")
for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]):
with self.session(graph=ops.Graph()):
# Initialize variables for numpy implementation.
@@ -171,8 +174,8 @@
var1 = resource_variable_ops.ResourceVariable(
var1_np, name="var1_%d" % i)
else:
- var0 = variables.Variable(var0_np)
- var1 = variables.Variable(var1_np)
+ var0 = variables.RefVariable(var0_np)
+ var1 = variables.RefVariable(var1_np)
grads0 = constant_op.constant(grads0_np)
grads1 = constant_op.constant(grads1_np)
@@ -194,6 +197,14 @@
self.assertTrue(beta2_power is not None)
self.assertIn(beta1_power, opt_variables)
self.assertIn(beta2_power, opt_variables)
+ # Ensure that non-slot variables are the same type as the requested
+ # variables.
+ self.assertEqual(
+ use_resource,
+ resource_variable_ops.is_resource_variable(beta1_power))
+ self.assertEqual(
+ use_resource,
+ resource_variable_ops.is_resource_variable(beta2_power))
if not context.executing_eagerly():
with ops.Graph().as_default():
diff --git a/tensorflow/python/training/basic_session_run_hooks.py b/tensorflow/python/training/basic_session_run_hooks.py
index b64c7ad..86718ab 100644
--- a/tensorflow/python/training/basic_session_run_hooks.py
+++ b/tensorflow/python/training/basic_session_run_hooks.py
@@ -163,7 +163,7 @@
return None
-@tf_export("train.LoggingTensorHook")
+@tf_export(v1=["train.LoggingTensorHook"])
class LoggingTensorHook(session_run_hook.SessionRunHook):
"""Prints the given tensors every N local steps, every N seconds, or at end.
@@ -373,7 +373,7 @@
self._update_steps_per_run_variable(global_step, run_context.session)
-@tf_export("train.StopAtStepHook")
+@tf_export(v1=["train.StopAtStepHook"])
class StopAtStepHook(session_run_hook.SessionRunHook):
"""Hook that requests stop at a specified step."""
@@ -495,7 +495,7 @@
pass
-@tf_export("train.CheckpointSaverHook")
+@tf_export(v1=["train.CheckpointSaverHook"])
class CheckpointSaverHook(session_run_hook.SessionRunHook):
"""Saves checkpoints every N steps or seconds."""
@@ -634,7 +634,7 @@
return savers[0]
-@tf_export("train.StepCounterHook")
+@tf_export(v1=["train.StepCounterHook"])
class StepCounterHook(session_run_hook.SessionRunHook):
"""Hook that counts steps per second."""
@@ -725,7 +725,7 @@
return "NaN loss during training."
-@tf_export("train.NanTensorHook")
+@tf_export(v1=["train.NanTensorHook"])
class NanTensorHook(session_run_hook.SessionRunHook):
"""Monitors the loss tensor and stops training if loss is NaN.
@@ -757,7 +757,7 @@
run_context.request_stop()
-@tf_export("train.SummarySaverHook")
+@tf_export(v1=["train.SummarySaverHook"])
class SummarySaverHook(session_run_hook.SessionRunHook):
"""Saves summaries every N steps."""
@@ -866,7 +866,7 @@
return summary_op
-@tf_export("train.GlobalStepWaiterHook")
+@tf_export(v1=["train.GlobalStepWaiterHook"])
class GlobalStepWaiterHook(session_run_hook.SessionRunHook):
"""Delays execution until global step reaches `wait_until_step`.
@@ -914,7 +914,7 @@
time.sleep(0.5)
-@tf_export("train.FinalOpsHook")
+@tf_export(v1=["train.FinalOpsHook"])
class FinalOpsHook(session_run_hook.SessionRunHook):
"""A hook which evaluates `Tensors` at the end of a session."""
@@ -958,7 +958,7 @@
raise e
-@tf_export("train.FeedFnHook")
+@tf_export(v1=["train.FeedFnHook"])
class FeedFnHook(session_run_hook.SessionRunHook):
"""Runs `feed_fn` and sets the `feed_dict` accordingly."""
diff --git a/tensorflow/python/training/basic_session_run_hooks_test.py b/tensorflow/python/training/basic_session_run_hooks_test.py
index 08942c5..1af2762 100644
--- a/tensorflow/python/training/basic_session_run_hooks_test.py
+++ b/tensorflow/python/training/basic_session_run_hooks_test.py
@@ -1122,6 +1122,7 @@
self.assertGreater(summary_value.simple_value, 0)
+@test_util.run_v1_only('b/120545219')
class SummarySaverHookTest(test.TestCase):
def setUp(self):
@@ -1148,13 +1149,11 @@
basic_session_run_hooks.SummarySaverHook(
scaffold=monitored_session.Scaffold(), summary_op=self.summary_op)
- @test_util.run_deprecated_v1
def test_raise_in_both_secs_and_steps(self):
with self.assertRaises(ValueError):
basic_session_run_hooks.SummarySaverHook(
save_secs=10, save_steps=20, summary_writer=self.summary_writer)
- @test_util.run_deprecated_v1
def test_raise_in_none_secs_and_steps(self):
with self.assertRaises(ValueError):
basic_session_run_hooks.SummarySaverHook(
@@ -1405,6 +1404,7 @@
hook.final_ops_values.tolist())
+@test_util.run_v1_only('b/120545219')
class ResourceSummarySaverHookTest(test.TestCase):
def setUp(self):
diff --git a/tensorflow/python/training/checkpoint_management.py b/tensorflow/python/training/checkpoint_management.py
index f745ab4..a7ad1f7 100644
--- a/tensorflow/python/training/checkpoint_management.py
+++ b/tensorflow/python/training/checkpoint_management.py
@@ -56,10 +56,6 @@
return os.path.join(save_dir, latest_filename)
-@deprecation.deprecated(
- date=None,
- instructions=("Use tf.train.CheckpointManager to manage checkpoints rather "
- "than editing the Checkpoint proto manually."))
@tf_export(v1=["train.generate_checkpoint_state_proto"])
def generate_checkpoint_state_proto(save_dir,
model_checkpoint_path,
diff --git a/tensorflow/python/training/checkpoint_ops_test.py b/tensorflow/python/training/checkpoint_ops_test.py
index 21ad3df..a0fd2dc 100644
--- a/tensorflow/python/training/checkpoint_ops_test.py
+++ b/tensorflow/python/training/checkpoint_ops_test.py
@@ -23,6 +23,7 @@
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.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import partitioned_variables
@@ -33,6 +34,7 @@
from tensorflow.python.training import saver as saver_lib
+@test_util.run_v1_only('b/120545219')
class LoadAndRemapWrappersTest(test.TestCase):
"""Tests for the functionality of the Python wrappers."""
@@ -152,7 +154,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
@@ -186,7 +188,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
@@ -224,7 +226,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
@@ -260,7 +262,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_matrix,
remapped_matrix.as_tensor().eval())
@@ -294,7 +296,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor().eval())
@@ -340,7 +342,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor().eval())
@@ -378,7 +380,7 @@
partitioner=partitioned_variables.fixed_size_partitioner(2))
with self.cached_session():
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose(expected_remapped_embeddings,
remapped_embeddings.as_tensor().eval())
diff --git a/tensorflow/python/training/checkpoint_utils.py b/tensorflow/python/training/checkpoint_utils.py
index 58166db..74b4617 100644
--- a/tensorflow/python/training/checkpoint_utils.py
+++ b/tensorflow/python/training/checkpoint_utils.py
@@ -21,6 +21,7 @@
import six
from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.framework import ops
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import resource_variable_ops
@@ -29,8 +30,7 @@
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import checkpoint_management
-from tensorflow.python.training import distribution_strategy_context
-from tensorflow.python.training import saver
+from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.util.tf_export import tf_export
@@ -311,10 +311,10 @@
restore_op = io_ops.restore_v2(
ckpt_file, [tensor_name], [slice_spec], [base_type], name=name)[0]
- names_to_saveables = saver.BaseSaverBuilder.OpListToDict([variable])
+ names_to_saveables = saveable_object_util.op_list_to_dict([variable])
saveable_objects = []
for name, op in names_to_saveables.items():
- for s in saver.BaseSaverBuilder.SaveableObjectsForOp(op, name):
+ for s in saveable_object_util.saveable_objects_for_op(op, name):
saveable_objects.append(s)
assert len(saveable_objects) == 1 # Should be only one variable.
diff --git a/tensorflow/python/training/checkpointable/BUILD b/tensorflow/python/training/checkpointable/BUILD
index 4ab5593..3201c75 100644
--- a/tensorflow/python/training/checkpointable/BUILD
+++ b/tensorflow/python/training/checkpointable/BUILD
@@ -25,9 +25,9 @@
"//tensorflow/python:framework_ops",
"//tensorflow/python:io_ops_gen",
"//tensorflow/python:platform",
- "//tensorflow/python:saveable_object",
"//tensorflow/python:util",
"//tensorflow/python/eager:context",
+ "//tensorflow/python/training/saving:saveable_object",
],
)
@@ -114,7 +114,6 @@
"//tensorflow/python:init_ops",
"//tensorflow/python:io_ops_gen",
"//tensorflow/python:pywrap_tensorflow",
- "//tensorflow/python:saveable_object",
"//tensorflow/python:saver",
"//tensorflow/python:session",
"//tensorflow/python:tensor_shape",
@@ -123,6 +122,9 @@
"//tensorflow/python:variables",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
+ "//tensorflow/python/training/saving:functional_saver",
+ "//tensorflow/python/training/saving:saveable_object",
+ "//tensorflow/python/training/saving:saveable_object_util",
],
)
@@ -157,6 +159,43 @@
"//tensorflow/python/eager:test",
"//tensorflow/python/keras:engine",
"//tensorflow/python/keras:layers",
+ "@absl_py//absl/testing:parameterized",
+ "@six_archive//:six",
+ ],
+)
+
+py_test(
+ name = "util_with_v1_optimizers_test",
+ srcs = ["util_with_v1_optimizers_test.py"],
+ srcs_version = "PY2AND3",
+ tags = ["notsan"], # b/74395663
+ deps = [
+ ":base",
+ ":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/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/checkpointable/base.py b/tensorflow/python/training/checkpointable/base.py
index 095a90d..3cd1c6f 100644
--- a/tensorflow/python/training/checkpointable/base.py
+++ b/tensorflow/python/training/checkpointable/base.py
@@ -25,7 +25,6 @@
import six
-from tensorflow.python import pywrap_tensorflow
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
@@ -34,7 +33,7 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_io_ops as io_ops
from tensorflow.python.platform import tf_logging as logging
-from tensorflow.python.training import saveable_object
+from tensorflow.python.training.saving import saveable_object
from tensorflow.python.util import nest
from tensorflow.python.util import serialization
from tensorflow.python.util import tf_decorator
@@ -374,41 +373,10 @@
eagerly.
"""
(restore_ops,
- named_saveables,
+ tensor_saveables,
python_saveables) = self._gather_ops_or_named_saveables()
-
- # Eagerly run restorations for Python state.
- reader = pywrap_tensorflow.NewCheckpointReader(
- self._checkpoint.save_path_string)
- for saveable in python_saveables:
- spec_names = [spec.name for spec in saveable.specs]
- saveable.python_restore(
- [reader.get_tensor(name) for name in spec_names])
-
- # If we have new SaveableObjects, extract and cache restore ops.
- if named_saveables:
- validated_saveables = (
- self._checkpoint.builder._ValidateAndSliceInputs(named_saveables)) # pylint: disable=protected-access
- validated_names = set(saveable.name for saveable in validated_saveables)
- if set(named_saveables.keys()) != validated_names:
- raise AssertionError(
- ("Saveable keys changed when validating. Got back %s, was "
- "expecting %s") % (named_saveables.keys(), validated_names))
- all_tensors = self._checkpoint.builder.bulk_restore(
- filename_tensor=self._checkpoint.save_path_tensor,
- saveables=validated_saveables, preferred_shard=-1,
- restore_sequentially=False)
- saveable_index = 0
- for saveable in validated_saveables:
- num_specs = len(saveable.specs)
- saveable_tensors = all_tensors[
- saveable_index:saveable_index + num_specs]
- saveable_index += num_specs
- restore_op = saveable.restore(saveable_tensors, restored_shapes=None)
- if not context.executing_eagerly():
- assert saveable.name not in self._checkpoint.restore_ops_by_name
- self._checkpoint.restore_ops_by_name[saveable.name] = restore_op
- restore_ops.append(restore_op)
+ restore_ops.extend(self._checkpoint.restore_saveables(
+ tensor_saveables, python_saveables))
return restore_ops
@property
diff --git a/tensorflow/python/training/checkpointable/data_structures_test.py b/tensorflow/python/training/checkpointable/data_structures_test.py
index 9cefd94..bcec6e0 100644
--- a/tensorflow/python/training/checkpointable/data_structures_test.py
+++ b/tensorflow/python/training/checkpointable/data_structures_test.py
@@ -73,6 +73,7 @@
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]))
@@ -105,6 +106,7 @@
self.assertIn(v, model.trainable_variables)
self.assertNotIn(v, model.non_trainable_variables)
+ @test_util.run_v1_only("b/120545219")
def testUpdatesForwarded(self):
with context.graph_mode():
model = HasList()
@@ -121,6 +123,7 @@
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])
@@ -295,6 +298,7 @@
class MappingTests(test.TestCase):
@test_util.run_in_graph_and_eager_modes
+ @test_util.run_v1_only("b/120545219")
def testTracking(self):
model = HasMapping()
output = model(array_ops.ones([32, 2]))
diff --git a/tensorflow/python/training/checkpointable/util.py b/tensorflow/python/training/checkpointable/util.py
index a1cc608..fde9194 100644
--- a/tensorflow/python/training/checkpointable/util.py
+++ b/tensorflow/python/training/checkpointable/util.py
@@ -39,12 +39,15 @@
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.training import checkpoint_management
-from tensorflow.python.training import optimizer as optimizer_lib
-from tensorflow.python.training import saveable_object as saveable_object_lib
-from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training import optimizer as optimizer_v1
+from tensorflow.python.training import saver as v1_saver_lib
from tensorflow.python.training.checkpointable import base
from tensorflow.python.training.checkpointable import data_structures
from tensorflow.python.training.checkpointable import tracking
+from tensorflow.python.training.saving import functional_saver
+from tensorflow.python.training.saving import saveable_object as saveable_object_lib
+from tensorflow.python.training.saving import saveable_object_util
+from tensorflow.python.util import compat
from tensorflow.python.util import deprecation
from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export
@@ -89,7 +92,6 @@
referenced every restore (e.g. for Python state); otherwise they would
create their own ops every restore.
"""
- self.builder = saver_lib.BulkSaverBuilder()
self.object_graph_proto = object_graph_proto
self.restore_uid = ops.uid()
# Maps from objects to lists of attributes which were in the checkpoint but
@@ -144,6 +146,57 @@
if self.new_restore_ops_callback:
self.new_restore_ops_callback(new_ops) # pylint: disable=not-callable
+ def restore_saveables(self, tensor_saveables, python_saveables):
+ """Run or build restore operations for SaveableObjects.
+
+ Args:
+ tensor_saveables: `SaveableObject`s which correspond to Tensors.
+ python_saveables: `PythonStateSaveable`s which correspond to Python
+ values.
+
+ Returns:
+ When graph building, a list of restore operations, either cached or newly
+ created, to restore `tensor_saveables`.
+ """
+ restore_ops = []
+ # Eagerly run restorations for Python state.
+ reader = pywrap_tensorflow.NewCheckpointReader(
+ self.save_path_string)
+ for saveable in python_saveables:
+ spec_names = [spec.name for spec in saveable.specs]
+ saveable.python_restore(
+ [reader.get_tensor(name) for name in spec_names])
+
+ # If we have new SaveableObjects, extract and cache restore ops.
+ if tensor_saveables:
+ validated_saveables = saveable_object_util.validate_and_slice_inputs(
+ tensor_saveables)
+ validated_names = set(saveable.name for saveable in validated_saveables)
+ if set(tensor_saveables.keys()) != validated_names:
+ raise AssertionError(
+ ("Saveable keys changed when validating. Got back %s, was "
+ "expecting %s") % (tensor_saveables.keys(), validated_names))
+ for saveable in validated_saveables:
+ if saveable.device:
+ device = saveable_object_util.set_cpu0(saveable.device)
+ else:
+ device = None
+ with ops.device(device):
+ tensors = []
+ for spec in saveable.specs:
+ tensors.append(
+ io_ops.restore_v2(
+ self.save_path_tensor,
+ [spec.name],
+ [spec.slice_spec],
+ [spec.dtype])[0])
+ restore_op = saveable.restore(tensors, restored_shapes=None)
+ if not context.executing_eagerly():
+ assert saveable.name not in self.restore_ops_by_name
+ self.restore_ops_by_name[saveable.name] = restore_op
+ restore_ops.append(restore_op)
+ return restore_ops
+
class _NameBasedRestoreCoordinator(object):
"""Keeps the status of a name-based checkpoint restore."""
@@ -183,11 +236,11 @@
continue
else:
saveable = saveable_factory
- names_to_saveables = saver_lib.BaseSaverBuilder.OpListToDict(
+ names_to_saveables = saveable_object_util.op_list_to_dict(
[saveable],
convert_variable_to_tensor=False)
for name, op in names_to_saveables.items():
- for saveable_object in saver_lib.BaseSaverBuilder.SaveableObjectsForOp(
+ for saveable_object in saveable_object_util.saveable_objects_for_op(
op=op, name=name):
yield saveable_object
@@ -507,7 +560,9 @@
non_slot_objects = list(checkpointable_objects)
slot_variables = _ObjectIdentityDictionary()
for checkpointable in non_slot_objects:
- if isinstance(checkpointable, optimizer_lib.Optimizer):
+ if (isinstance(checkpointable, optimizer_v1.Optimizer)
+ # TODO(b/110718070): Fix Keras imports.
+ or hasattr(checkpointable, "_create_or_restore_slot_variable")):
naming_scheme = _slot_variable_naming_for_optimizer(
optimizer_path=object_names[checkpointable])
slot_names = checkpointable.get_slot_names()
@@ -517,7 +572,7 @@
try:
slot_variable = checkpointable.get_slot(
original_variable, slot_name)
- except AttributeError:
+ except (AttributeError, KeyError):
slot_variable = None
if slot_variable is None:
continue
@@ -606,10 +661,10 @@
# Figure out the name-based Saver's name for this variable. If it's
# already a SaveableObject we'd just get the checkpoint key back, so
# we leave full_name blank.
- saver_dict = saver_lib.BaseSaverBuilder.OpListToDict(
+ saver_dict = saveable_object_util.op_list_to_dict(
[maybe_saveable], convert_variable_to_tensor=False)
full_name, = saver_dict.keys()
- saveables = tuple(saver_lib.BaseSaverBuilder.SaveableObjectsForOp(
+ saveables = tuple(saveable_object_util.saveable_objects_for_op(
op=maybe_saveable, name=attribute.checkpoint_key))
for saveable in saveables:
saveable.full_name = full_name
@@ -649,8 +704,10 @@
return named_saveable_objects, feed_additions
-def _fill_object_graph_proto(checkpointable_objects, node_ids, slot_variables,
- object_graph_proto=None):
+def fill_object_graph_proto(checkpointable_objects,
+ node_ids,
+ slot_variables,
+ object_graph_proto=None):
"""Name non-slot `Checkpointable`s and add them to `object_graph_proto`."""
if object_graph_proto is None:
object_graph_proto = (
@@ -679,7 +736,7 @@
checkpointable_objects=checkpointable_objects,
node_ids=node_ids,
object_names=object_names)
- object_graph_proto = _fill_object_graph_proto(
+ object_graph_proto = fill_object_graph_proto(
checkpointable_objects=checkpointable_objects,
node_ids=node_ids,
slot_variables=slot_variables)
@@ -732,7 +789,7 @@
return _serialize_object_graph(root_checkpointable, None)[0]
-def _find_objects(root_checkpointable):
+def find_objects(root_checkpointable):
"""Find and number objects which are dependencies of `root_checkpointable`."""
checkpointable_objects, path_to_root = (
_breadth_first_checkpointable_traversal(root_checkpointable))
@@ -763,18 +820,10 @@
Returns:
A flat list of objects.
"""
- checkpointable_objects, _, _ = _find_objects(root_checkpointable)
+ checkpointable_objects, _, _ = find_objects(root_checkpointable)
return checkpointable_objects
-def make_object_graph_without_attributes(root_checkpointable, proto=None):
- """Fill an object graph proto, ignoring variable values."""
- checkpointable_objects, node_ids, slot_variables = _find_objects(
- root_checkpointable)
- return _fill_object_graph_proto(
- checkpointable_objects, node_ids, slot_variables, proto)
-
-
def gather_initializers(root_checkpointable):
"""Traverse the object graph and find initialization ops.
@@ -1232,7 +1281,7 @@
session = ops.get_default_session()
with ops.device("/cpu:0"):
saveables = self._gather_saveable_objects()
- saver_lib.Saver(saveables).restore(
+ v1_saver_lib.Saver(saveables).restore(
sess=session, save_path=self._checkpoint.save_path)
def initialize_or_restore(self, session=None):
@@ -1257,18 +1306,6 @@
fetches=fetches, feed_dict=feed_dict, **kwargs)
-def _copy_saver_with_new_var_list(old_saver, new_var_list):
- """Copy a `tf.train.Saver`'s state to a new Saver with different variables."""
- new_saver = saver_lib.Saver(var_list=new_var_list, max_to_keep=None)
- # TODO(allenl): Move to copying functionality to Saver?
- # pylint: disable=protected-access
- new_saver._last_checkpoints = old_saver._last_checkpoints
- new_saver._checkpoints_to_be_deleted = old_saver._checkpoints_to_be_deleted
- new_saver._next_checkpoint_time = old_saver._next_checkpoint_time
- # pylint: enable=protected-access
- return new_saver
-
-
class CheckpointableSaver(object):
"""Saves and restores a `Checkpointable` object and its dependencies.
@@ -1307,7 +1344,8 @@
# Op caching for save
self._object_graph_feed_tensor = None
self._last_save_object_graph = None
- self._last_save_saver = None
+ self._file_prefix_feed_tensor = None
+ self._cached_save_operation = None
# Op caching for restore, shared between _CheckpointRestoreCoordinators
self._restore_op_cache = {}
@@ -1374,13 +1412,16 @@
base.NoRestoreSaveable(
tensor=object_graph_tensor,
name=base.OBJECT_GRAPH_PROTO_KEY))
- # TODO(allenl, haoliang): Swap in a function-based saver here.
- return saver_lib.Saver(
+ # TODO(allenl): Swap in a function-based saver here once it can serialize
+ # to a SaverDef.
+ return v1_saver_lib.Saver(
var_list=named_saveable_objects, max_to_keep=None)
- def _prepare_save(self,
- object_graph_tensor=None,
- saveable_object_cache=None):
+ def _save_cached_when_graph_building(
+ self,
+ file_prefix,
+ object_graph_tensor=None,
+ saveable_object_cache=None):
"""Create or retrieve save ops.
When graph building, `saveable_object_cache` will typically be non-`None`,
@@ -1389,15 +1430,17 @@
unnecessarily re-creating save ops.
Args:
+ file_prefix: The prefix for saved checkpoint files.
object_graph_tensor: A `Tensor` to which the current object graph will be
fed.
saveable_object_cache: A dictionary; if specified, used to cache
`SaveableObject`s.
Returns:
- A two-element tuple with a `tf.train.Saver` and a feed_dict of `Tensor`s
- to feed when running save ops. The feed dict contains the current object
- graph and any Python state to be saved in the checkpoint.
+ A two-element tuple with a filename tensor and a feed_dict of tensors to
+ feed when running it (if graph building). The feed dict contains the
+ current object graph and any Python state to be saved in the
+ checkpoint. When executing eagerly only the first argument is meaningful.
"""
(named_saveable_objects, graph_proto,
feed_additions) = self._gather_saveables(
@@ -1409,15 +1452,11 @@
# constructors. That means the Saver needs to be copied with a new
# var_list.
or context.executing_eagerly()):
- if self._last_save_object_graph is not None:
- self._last_save_saver = _copy_saver_with_new_var_list(
- old_saver=self._last_save_saver,
- new_var_list=named_saveable_objects)
- else:
- self._last_save_saver = saver_lib.Saver(
- var_list=named_saveable_objects, max_to_keep=None)
+ saver = functional_saver.Saver(named_saveable_objects)
+ with ops.device("/cpu:0"):
+ self._cached_save_operation = saver.save(file_prefix)
self._last_save_object_graph = graph_proto
- return self._last_save_saver, feed_additions
+ return self._cached_save_operation, feed_additions
def save(self, file_prefix, checkpoint_number=None, session=None):
"""Save a training checkpoint.
@@ -1441,36 +1480,42 @@
Returns:
The full path to the checkpoint.
"""
- feed_additions = {}
+ feed_dict = {}
graph_building = not context.executing_eagerly()
+ if checkpoint_number:
+ file_prefix = "%s-%d" % (file_prefix, checkpoint_number)
if graph_building:
if self._object_graph_feed_tensor is None:
with ops.device("/cpu:0"):
self._object_graph_feed_tensor = constant_op.constant(
"", dtype=dtypes.string)
+ self._file_prefix_feed_tensor = constant_op.constant(
+ "", dtype=dtypes.string)
object_graph_tensor = self._object_graph_feed_tensor
+ file_prefix_tensor = self._file_prefix_feed_tensor
+ feed_dict[file_prefix_tensor] = file_prefix
else:
+ with ops.device("/cpu:0"):
+ file_prefix_tensor = constant_op.constant(
+ file_prefix, dtype=dtypes.string)
object_graph_tensor = None
- saver, new_feed_additions = self._prepare_save(
+ file_io.recursive_create_dir(os.path.dirname(file_prefix))
+ save_path, new_feed_additions = self._save_cached_when_graph_building(
+ file_prefix=file_prefix_tensor,
object_graph_tensor=object_graph_tensor,
saveable_object_cache=self._saveable_object_cache)
if new_feed_additions:
- feed_additions.update(new_feed_additions)
+ feed_dict.update(new_feed_additions)
if not graph_building:
session = None
elif session is None:
session = ops.get_default_session()
- file_io.recursive_create_dir(os.path.dirname(file_prefix))
- with ops.device("/cpu:0"):
- save_path = saver.save(
- sess=_SessionWithFeedDictAdditions(
- session=session, feed_additions=feed_additions),
- save_path=file_prefix,
- write_meta_graph=False,
- write_state=False,
- global_step=checkpoint_number)
+ if session:
+ save_path = session.run(save_path, feed_dict=feed_dict)
+ else:
+ save_path = save_path.numpy()
return save_path
def restore(self, save_path):
@@ -1759,9 +1804,9 @@
Returns:
The full path to the checkpoint (i.e. `file_prefix`).
"""
- return self._saver.save(
+ return compat.as_str(self._saver.save(
file_prefix=file_prefix,
- session=session)
+ session=session))
@property
def save_counter(self):
@@ -1820,7 +1865,7 @@
checkpoint_number = assign_op.numpy()
file_path = self.write("%s-%d" % (file_prefix, checkpoint_number),
session=session)
- checkpoint_management.update_checkpoint_state(
+ checkpoint_management.update_checkpoint_state_internal(
save_dir=os.path.dirname(file_prefix),
model_checkpoint_path=file_path,
all_model_checkpoint_paths=[file_path])
diff --git a/tensorflow/python/training/checkpointable/util_test.py b/tensorflow/python/training/checkpointable/util_test.py
index de9cac0..61de468 100644
--- a/tensorflow/python/training/checkpointable/util_test.py
+++ b/tensorflow/python/training/checkpointable/util_test.py
@@ -20,10 +20,10 @@
import json
import os
+from absl.testing import parameterized
import six
from tensorflow.python import pywrap_tensorflow
-from tensorflow.python.client import session as session_lib
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
@@ -35,16 +35,15 @@
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
-from tensorflow.python.training import adam
+from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.training import checkpoint_management
-from tensorflow.python.training import momentum
from tensorflow.python.training import saver as saver_lib
from tensorflow.python.training import training_util
from tensorflow.python.training.checkpointable import base
@@ -199,17 +198,6 @@
with self.assertRaises(NotImplementedError):
checkpoint_reversed.save(prefix)
- @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True)
- def test_object_graph_no_attributes(self):
- root = tracking.Checkpointable()
- root.v = resource_variable_ops.ResourceVariable(1.)
- root.opt = momentum.MomentumOptimizer(0.01, 0.5)
- root.opt.minimize(root.v.read_value)
- object_graph = checkpointable_utils.make_object_graph_without_attributes(
- root)
- # Four objects: Root, v, opt, and a slot variable for v
- self.assertEqual(4, len(object_graph.nodes))
-
class _MirroringSaveable(saver_lib.BaseSaverBuilder.SaveableObject):
@@ -255,7 +243,7 @@
return self.non_dep_variable.name
-class CheckpointingTests(test.TestCase):
+class CheckpointingTests(parameterized.TestCase, test.TestCase):
@test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True)
def testNamingWithOptimizer(self):
@@ -264,41 +252,32 @@
# 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()
+ optimizer = adam.Adam(0.001)
+ step = training_util.get_or_create_global_step()
root_checkpointable = checkpointable_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(checkpointable_utils.gather_initializers(
- root_checkpointable))
- self.evaluate(train_op)
+ 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(checkpointable_utils.gather_initializers(
+ root_checkpointable))
+ self.evaluate(train_op)
named_variables, serialized_graph, _ = (
checkpointable_utils._serialize_object_graph(
root_checkpointable, saveables_cache=None))
- 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
+ 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",
@@ -306,9 +285,26 @@
"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/epsilon",
+ "optimizer/iter",
+ "optimizer/decay",
+ ) + expected_slot_keys
suffix = "/.ATTRIBUTES/VARIABLE_VALUE"
expected_checkpoint_names = [
name + suffix for name in expected_checkpoint_names]
+ expected_checkpoint_names.append(
+ "optimizer/.ATTRIBUTES/OBJECT_CONFIG_JSON")
# The Dense layers also save get_config() JSON
expected_checkpoint_names.extend(
["model/_second/.ATTRIBUTES/OBJECT_CONFIG_JSON",
@@ -319,7 +315,7 @@
# Check that we've mapped to the right variable objects (not exhaustive)
self.assertEqual(
"global_step",
- named_variables["optimizer_step" + suffix].full_name)
+ named_variables["step" + suffix].full_name)
self.assertEqual(
"my_model/dense_1/kernel",
named_variables["model/_second/kernel" + suffix].full_name)
@@ -327,48 +323,31 @@
"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)
+ "beta_1",
+ named_variables["optimizer/beta_1" + suffix].full_name)
self.assertEqual(
- "beta2_power",
- named_variables["optimizer/beta2_power" + suffix].full_name)
+ "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]
- 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)
+ children = [node.local_name for node in optimizer_node.children]
+ six.assertCountEqual(
+ self,
+ # Non-slot dependencies
+ ["beta_1", "beta_2", "iter", "decay", "epsilon", "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):
@@ -409,20 +388,19 @@
@test_util.run_in_graph_and_eager_modes
def testSaveRestore(self):
model = MyModel()
- optimizer = adam.AdamOptimizer(0.001)
+ optimizer = adam.Adam(0.001)
root_checkpointable = checkpointable_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_checkpointable.save_counter # pylint: disable=pointless-statement
- self.evaluate(checkpointable_utils.gather_initializers(
- root_checkpointable))
- self.evaluate(train_op)
+ 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))
+ root_checkpointable.save_counter # pylint: disable=pointless-statement
+ self.evaluate(checkpointable_utils.gather_initializers(
+ root_checkpointable))
+ 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")
@@ -430,7 +408,8 @@
save_path = root_checkpointable.save(file_prefix=prefix)
self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.]))
self.evaluate(state_ops.assign(root_checkpointable.save_counter, 3))
- optimizer_variables = self.evaluate(optimizer.variables())
+ 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_checkpointable.restore(save_path=save_path).assert_consumed()
@@ -441,11 +420,7 @@
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 appying gradients so we can
- # test that they've been restored correctly.
- beta1=1.0, beta2=1.0)
+ on_create_optimizer = adam.Adam(0.001)
on_create_root = checkpointable_utils.Checkpoint(
optimizer=on_create_optimizer, model=on_create_model)
# Deferred restoration
@@ -467,15 +442,15 @@
# 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)
+ on_create_optimizer.minimize(loss=dummy_var.read_value,
+ var_list=[dummy_var])
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))
+ 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):
@@ -485,21 +460,22 @@
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
for training_continuation in range(3):
model = MyModel()
- optimizer = adam.AdamOptimizer(0.001)
+ optimizer = adam.Adam(0.001)
root = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model,
- optimizer_step=training_util.get_or_create_global_step())
+ 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.]])
- optimizer.minimize(
- lambda: model(input_value), # pylint: disable=cell-var-from-loop
- global_step=root.optimizer_step)
+ 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_step.numpy())
+ root.optimizer.iterations.numpy())
def testUsageGraph(self):
"""Expected usage when graph building."""
@@ -510,14 +486,16 @@
for training_continuation in range(3):
with ops.Graph().as_default():
model = MyModel()
- optimizer = adam.AdamOptimizer(0.001)
+ optimizer = adam.Adam(0.001)
root = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model,
- global_step=training_util.get_or_create_global_step())
+ optimizer=optimizer, model=model)
input_value = constant_op.constant([[3.]])
- train_op = optimizer.minimize(
- model(input_value),
- global_step=root.global_step)
+ 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:
@@ -536,7 +514,7 @@
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))
+ session.run(root.optimizer.iterations))
self.assertEqual(training_continuation + 1,
session.run(root.save_counter))
@@ -546,21 +524,23 @@
# 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.AdamOptimizer(0.001)
+ optimizer = adam.Adam(0.001)
root = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model,
- global_step=training_util.get_or_create_global_step())
+ 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(
- optimizer.minimize,
- functools.partial(model, input_value),
- global_step=root.global_step)
+ 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()
@@ -568,7 +548,7 @@
train_fn()
manager.save()
self.assertEqual((training_continuation + 1) * num_training_steps,
- self.evaluate(root.global_step))
+ self.evaluate(root.optimizer.iterations))
self.assertEqual(training_continuation + 1,
self.evaluate(root.save_counter))
@@ -628,6 +608,7 @@
# 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()
@@ -636,10 +617,9 @@
with test_util.device(use_gpu=True):
model = MyModel()
# Don't actually train so we can test variable values
- optimizer = adam.AdamOptimizer(0.)
+ optimizer = adam.Adam(0.)
root = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model,
- global_step=training_util.get_or_create_global_step())
+ optimizer=optimizer, model=model)
checkpoint_path = checkpoint_management.latest_checkpoint(
checkpoint_directory)
status = root.restore(save_path=checkpoint_path)
@@ -650,8 +630,7 @@
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)
+ return optimizer.apply_gradients(zip(gradients, model.variables))
if not context.executing_eagerly():
train_fn = functools.partial(
self.evaluate, train_fn())
@@ -665,7 +644,7 @@
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.evaluate(optimizer.iterations))
self.assertEqual(training_continuation + 1,
self.evaluate(root.save_counter))
# pylint: enable=cell-var-from-loop
@@ -727,7 +706,7 @@
with context.eager_mode():
model = Model()
- optimizer = adam.AdamOptimizer(learning_rate=0.05)
+ optimizer = adam.Adam(learning_rate=0.05)
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
checkpoint = checkpointable_utils.Checkpoint(
@@ -813,24 +792,24 @@
root = tracking.Checkpointable()
root.var = checkpointable_utils.add_variable(
root, name="var", initializer=0.)
- optimizer = adam.AdamOptimizer(0.1)
- if context.executing_eagerly():
- optimizer.minimize(root.var.read_value)
- else:
- train_op = optimizer.minimize(root.var)
- # 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(checkpointable_utils.gather_initializers(
- checkpointable_utils.Checkpoint(root=root, optimizer=optimizer)))
- self.evaluate(train_op)
+ 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(checkpointable_utils.gather_initializers(
+ checkpointable_utils.Checkpoint(root=root, optimizer=optimizer)))
+ self.evaluate(train_op)
self.evaluate(state_ops.assign(root.var, 12.))
no_slots_path = checkpointable_utils.CheckpointableSaver(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(name="m", var=root.var),
- 14.))
+ self.evaluate(state_ops.assign(
+ optimizer.get_slot(slot_name="m", var=root.var),
+ 14.))
slots_path = checkpointable_utils.CheckpointableSaver(root).save(
os.path.join(checkpoint_directory, "with_slots"))
new_root = tracking.Checkpointable()
@@ -847,29 +826,32 @@
no_slot_status.assert_consumed()
no_slot_status.run_restore_ops()
self.assertEqual(12., self.evaluate(new_root.var))
- new_root.optimizer = adam.AdamOptimizer(0.1)
+ new_root.optimizer = adam.Adam(0.1)
slot_status.assert_existing_objects_matched()
- with self.assertRaisesRegexp(AssertionError, "beta1_power"):
+ 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(name="m", var=new_root.var)))
+ new_root.optimizer.get_slot(slot_name="m", var=new_root.var)))
else:
- self.assertIs(new_root.optimizer.get_slot(name="m", var=new_root.var),
- None)
- if context.executing_eagerly():
- new_root.optimizer.minimize(new_root.var.read_value)
- else:
- train_op = new_root.optimizer.minimize(new_root.var)
- # 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()
+ # 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(name="m", var=new_root.var)))
- self.evaluate(train_op)
+ 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
@@ -1029,18 +1011,18 @@
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
first = tracking.Checkpointable()
- first.var1 = variables.Variable(0., name="outside_var")
- first.var2 = variables.Variable(0., name="blah")
+ first.var1 = variables_lib.Variable(0., name="outside_var")
+ first.var2 = variables_lib.Variable(0., name="blah")
self.evaluate(first.var1.assign(4.))
self.evaluate(first.var2.assign(8.))
save_path = checkpointable_utils.CheckpointableSaver(first).save(
checkpoint_prefix)
second = tracking.Checkpointable()
- second.var2 = variables.Variable(0., name="blah")
+ second.var2 = variables_lib.Variable(0., name="blah")
status = checkpointable_utils.CheckpointableSaver(
second).restore(save_path)
- recreated_var1 = variables.Variable(0., name="outside_var")
+ recreated_var1 = variables_lib.Variable(0., name="outside_var")
status.run_restore_ops()
self.assertEqual(8., self.evaluate(second.var2))
self.evaluate(recreated_var1.assign(-2.))
@@ -1057,15 +1039,16 @@
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
obj = tracking.Checkpointable()
- obj.var = variable_scope.get_variable(name="v", initializer=0.)
- obj.opt = adam.AdamOptimizer(0.1)
- obj.opt.minimize(obj.var.read_value())
+ 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(checkpointable_utils.gather_initializers(obj))
saver = checkpointable_utils.CheckpointableSaver(obj)
saver.save(checkpoint_prefix)
- before_ops = graph.get_operations()
+ graph.finalize()
saver.save(checkpoint_prefix)
- self.assertEqual(before_ops, graph.get_operations())
@test_util.run_in_graph_and_eager_modes
def testCheckpointState(self):
@@ -1146,74 +1129,17 @@
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
obj = tracking.Checkpointable()
- obj.var = variable_scope.get_variable(name="v", initializer=0.)
- obj.opt = adam.AdamOptimizer(0.1)
- obj.opt.minimize(obj.var.read_value())
+ 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(checkpointable_utils.gather_initializers(obj))
saver = checkpointable_utils.CheckpointableSaver(obj)
save_path = saver.save(checkpoint_prefix)
saver.restore(save_path)
- before_ops = graph.get_operations()
+ graph.finalize()
saver.restore(save_path)
- self.assertEqual(before_ops, graph.get_operations())
-
- def testMultipleGraphsNonSlotVariables(self):
- with context.graph_mode():
- checkpoint_directory = self.get_temp_dir()
- checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
- optimizer = adam.AdamOptimizer(0.001)
- # Construct a model in one graph
- first_graph = ops.Graph()
- first_session = session_lib.Session(graph=first_graph)
- with first_graph.as_default(), first_session.as_default():
- first_variable = resource_variable_ops.ResourceVariable([1.])
- first_root_checkpointable = checkpointable_utils.Checkpoint(
- optimizer=optimizer, variable=first_variable)
- train_op = optimizer.minimize(first_variable.read_value)
- self.evaluate(checkpointable_utils.gather_initializers(
- first_root_checkpointable))
- self.evaluate(train_op)
- self.evaluate(first_variable.assign([1.]))
- self.evaluate(optimizer.get_slot(
- var=first_variable, name="m").assign([2.]))
- beta1_power, _ = optimizer._get_beta_accumulators()
- self.evaluate(beta1_power.assign(3.))
-
- # Save and load in a second graph
- second_graph = ops.Graph()
- with second_graph.as_default(), session_lib.Session(graph=second_graph):
- second_variable = resource_variable_ops.ResourceVariable([1.])
- second_root_checkpointable = checkpointable_utils.Checkpoint(
- optimizer=optimizer, variable=second_variable)
- train_op = optimizer.minimize(second_variable.read_value)
- second_root_checkpointable.restore(None).initialize_or_restore()
- self.evaluate(train_op)
- self.evaluate(second_variable.assign([4.]))
- self.evaluate(optimizer.get_slot(
- var=second_variable, name="m").assign([5.]))
- beta1_power, _ = optimizer._get_beta_accumulators()
- self.evaluate(beta1_power.assign(6.))
- save_path = second_root_checkpointable.save(checkpoint_prefix)
- self.evaluate(second_variable.assign([7.]))
- self.evaluate(optimizer.get_slot(
- var=second_variable, name="m").assign([8.]))
- beta1_power, _ = optimizer._get_beta_accumulators()
- self.assertAllEqual(6., self.evaluate(beta1_power))
- status = second_root_checkpointable.restore(save_path)
- status.assert_consumed().run_restore_ops()
- self.assertAllEqual([4.], self.evaluate(second_variable))
- self.assertAllEqual([5.], self.evaluate(optimizer.get_slot(
- var=second_variable, name="m")))
- beta1_power, _ = optimizer._get_beta_accumulators()
- self.assertAllEqual(6., self.evaluate(beta1_power))
-
- # Check that the first graph is unmolested
- with first_graph.as_default(), first_session.as_default():
- self.assertAllEqual([1.], self.evaluate(first_variable))
- self.assertAllEqual([2.], self.evaluate(optimizer.get_slot(
- var=first_variable, name="m")))
- beta1_power, _ = optimizer._get_beta_accumulators()
- self.assertAllEqual(3., self.evaluate(beta1_power))
@test_util.run_in_graph_and_eager_modes
def test_sequential(self):
@@ -1254,10 +1180,9 @@
optimizer_only_prefix = os.path.join(checkpoint_directory, "opt")
with test_util.device(use_gpu=True):
model = MyModel()
- optimizer = adam.AdamOptimizer(0.001)
+ optimizer = adam.Adam(0.001)
root = checkpointable_utils.Checkpoint(
- model=model, # Do not save the optimizer with the checkpoint.
- global_step=training_util.get_or_create_global_step())
+ model=model) # Do not save the optimizer with the checkpoint.
optimizer_checkpoint = checkpointable_utils.Checkpoint(
optimizer=optimizer)
@@ -1265,65 +1190,78 @@
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)
+ 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()
- self.evaluate([v.initializer for v in optimizer.variables()])
+ # 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.variables()[0].assign(42.))
+ 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.AdamOptimizer(0.001)
+ optimizer = adam.Adam(0.001)
root = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model,
- global_step=training_util.get_or_create_global_step())
+ optimizer=optimizer, model=model)
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)
+ 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_fn = functools.partial(self.evaluate, train_fn())
+ train_fn1 = functools.partial(self.evaluate, train_fn1())
status.initialize_or_restore()
- train_fn()
+ 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.AdamOptimizer(0.001, beta1=1.0)
+ optimizer = adam.Adam(0.001, beta1=1.0)
root = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model,
- global_step=training_util.get_or_create_global_step())
+ optimizer=optimizer, model=model)
opt_root = checkpointable_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)
+ 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_fn = functools.partial(self.evaluate, train_fn())
+ 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_fn()
- self.assertEqual(42., self.evaluate(optimizer.variables()[0]))
+ train_fn2()
+ self.assertEqual(42., self.evaluate(optimizer.beta_1))
@test_util.run_in_graph_and_eager_modes
def test_restore_after_adding_empty_checkpointable_data_structure(self):
@@ -1356,7 +1294,7 @@
return variable_scope.get_variable(name="in_manual_scope", shape=[])
-class TemplateTests(test.TestCase):
+class TemplateTests(parameterized.TestCase, test.TestCase):
@test_util.run_in_graph_and_eager_modes
def test_checkpointable_save_restore(self):
@@ -1380,10 +1318,11 @@
manual_dep, = manual_scope._checkpoint_dependencies
self.assertEqual("in_manual_scope", manual_dep.name)
self.assertIs(manual_scope_v, manual_dep.ref)
- optimizer = adam.AdamOptimizer(0.0)
+ optimizer = adam.Adam(0.0)
save_root = checkpointable_utils.Checkpoint(
my_template=save_template, optimizer=optimizer)
- optimizer.minimize(v1_save.read_value)
+ optimizer.minimize(v1_save.read_value,
+ var_list=[v1_save])
self.evaluate([v.initializer for v in save_template.variables])
self.evaluate([v.initializer for v in optimizer.variables()])
self.evaluate(v1_save.assign([12.]))
@@ -1393,13 +1332,13 @@
save_path = save_root.save(checkpoint_prefix)
load_template = template.make_template("s2", _templated)
- load_optimizer = adam.AdamOptimizer(0.0)
+ load_optimizer = adam.Adam(0.0)
load_root = checkpointable_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)
- self.assertEqual(3, len(load_template._checkpoint_dependencies))
+ 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",
@@ -1440,14 +1379,14 @@
status = load_root.restore(save_path)
(inner_template_one, inner_template_two), (v1, v2, v3) = load_template()
outer_template_dependencies = load_root.my_template._checkpoint_dependencies
- self.assertEqual(2, len(outer_template_dependencies))
+ self.assertLen(outer_template_dependencies, 2)
self.assertEqual("i1", outer_template_dependencies[0].name)
self.assertIs(inner_template_one, outer_template_dependencies[0].ref)
self.assertEqual("i2", outer_template_dependencies[1].name)
self.assertIs(inner_template_two, outer_template_dependencies[1].ref)
- self.assertEqual(1, len(inner_template_one._checkpoint_dependencies))
+ self.assertLen(inner_template_one._checkpoint_dependencies, 1)
self.assertEqual("v", inner_template_one._checkpoint_dependencies[0].name)
- self.assertEqual(1, len(inner_template_two._checkpoint_dependencies))
+ self.assertLen(inner_template_two._checkpoint_dependencies, 1)
self.assertEqual("v", inner_template_two._checkpoint_dependencies[0].name)
status.assert_consumed().run_restore_ops()
self.assertAllEqual([20.], self.evaluate(v1))
@@ -1460,13 +1399,14 @@
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()
+ optimizer = adam.Adam(0.001)
root_checkpointable = checkpointable_utils.Checkpoint(
- optimizer=optimizer, model=model, optimizer_step=optimizer_step)
- train_op = optimizer.minimize(
- functools.partial(model, input_value),
- global_step=optimizer_step)
+ 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(checkpointable_utils.gather_initializers(
root_checkpointable))
self.evaluate(train_op)
@@ -1474,28 +1414,26 @@
# 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.))
+ var=model._named_dense.bias, slot_name="m").assign([2.]))
+ self.evaluate(optimizer.beta_1.assign(3.))
return root_checkpointable
def _set_sentinels(self, root_checkpointable):
self.evaluate(root_checkpointable.model._named_dense.bias.assign([101.]))
self.evaluate(
root_checkpointable.optimizer.get_slot(
- var=root_checkpointable.model._named_dense.bias, name="m")
+ var=root_checkpointable.model._named_dense.bias, slot_name="m")
.assign([102.]))
- beta1_power, _ = root_checkpointable.optimizer._get_beta_accumulators()
- self.evaluate(beta1_power.assign(103.))
+ self.evaluate(root_checkpointable.optimizer.beta_1.assign(103.))
def _check_sentinels(self, root_checkpointable):
self.assertAllEqual(
[1.], self.evaluate(root_checkpointable.model._named_dense.bias))
self.assertAllEqual([2.], self.evaluate(
root_checkpointable.optimizer.get_slot(
- var=root_checkpointable.model._named_dense.bias, name="m")))
- beta1_power, _ = root_checkpointable.optimizer._get_beta_accumulators()
- self.assertAllEqual(3., self.evaluate(beta1_power))
+ var=root_checkpointable.model._named_dense.bias, slot_name="m")))
+ self.assertAllEqual(3.,
+ self.evaluate(root_checkpointable.optimizer.beta_1))
def _write_name_based_checkpoint(self):
checkpoint_directory = self.get_temp_dir()
@@ -1508,7 +1446,7 @@
name_saver = saver_lib.Saver()
return name_saver.save(
sess=session, save_path=checkpoint_prefix,
- global_step=root.optimizer_step)
+ global_step=root.optimizer.iterations)
@test_util.run_in_graph_and_eager_modes
def testLoadFromNameBasedSaver(self):
diff --git a/tensorflow/python/training/checkpointable/util_with_v1_optimizers_test.py b/tensorflow/python/training/checkpointable/util_with_v1_optimizers_test.py
new file mode 100644
index 0000000..00d5747
--- /dev/null
+++ b/tensorflow/python/training/checkpointable/util_with_v1_optimizers_test.py
@@ -0,0 +1,873 @@
+# 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.client import session as session_lib
+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.checkpointable import tracking
+from tensorflow.python.training.checkpointable import util as checkpointable_utils
+
+
+class NonLayerCheckpointable(tracking.Checkpointable):
+
+ def __init__(self):
+ super(NonLayerCheckpointable, self).__init__()
+ self.a_variable = checkpointable_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 Checkpointables which aren't Layers.
+ self._non_layer = NonLayerCheckpointable()
+
+ 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_checkpointable = checkpointable_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(checkpointable_utils.gather_initializers(
+ root_checkpointable))
+ self.evaluate(train_op)
+ named_variables, serialized_graph, _ = (
+ checkpointable_utils._serialize_object_graph(
+ root_checkpointable, saveables_cache=None))
+ 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]
+ # The Dense layers also save get_config() JSON
+ expected_checkpoint_names.extend(
+ ["model/_second/.ATTRIBUTES/OBJECT_CONFIG_JSON",
+ "model/_named_dense/.ATTRIBUTES/OBJECT_CONFIG_JSON"])
+ 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_checkpointable = checkpointable_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_checkpointable.save_counter # pylint: disable=pointless-statement
+ self.evaluate(checkpointable_utils.gather_initializers(
+ root_checkpointable))
+ 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_checkpointable.save(file_prefix=prefix)
+ self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.]))
+ self.evaluate(state_ops.assign(root_checkpointable.save_counter, 3))
+ optimizer_variables = self.evaluate(optimizer.variables())
+ self.evaluate(state_ops.assign(m_bias_slot, [-2.]))
+ # Immediate restoration
+ status = root_checkpointable.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_checkpointable.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 appying gradients so we can
+ # test that they've been restored correctly.
+ beta1=1.0, beta2=1.0)
+ on_create_root = checkpointable_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 = checkpointable_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 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 = checkpointable_utils.Checkpoint(
+ 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 = checkpointable_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 = checkpointable_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.Checkpointable()
+ checkpointable_utils.add_variable(
+ root, name=name, shape=[1, 2], dtype=dtypes.float64)
+ (named_variable,), _, _ = checkpointable_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 = checkpointable_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 = tracking.Checkpointable()
+ root.var = checkpointable_utils.add_variable(
+ root, name="var", initializer=0.)
+ optimizer = adam.AdamOptimizer(0.1)
+ if context.executing_eagerly():
+ optimizer.minimize(root.var.read_value)
+ else:
+ train_op = optimizer.minimize(root.var)
+ # 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(checkpointable_utils.gather_initializers(
+ checkpointable_utils.Checkpoint(root=root, optimizer=optimizer)))
+ self.evaluate(train_op)
+ self.evaluate(state_ops.assign(root.var, 12.))
+ no_slots_path = checkpointable_utils.CheckpointableSaver(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(name="m", var=root.var),
+ 14.))
+ slots_path = checkpointable_utils.CheckpointableSaver(root).save(
+ os.path.join(checkpoint_directory, "with_slots"))
+ new_root = tracking.Checkpointable()
+ # Load the slot-containing checkpoint (deferred), then immediately overwrite
+ # the non-slot variable (also deferred).
+ slot_status = checkpointable_utils.CheckpointableSaver(
+ new_root).restore(slots_path)
+ no_slot_status = checkpointable_utils.CheckpointableSaver(
+ new_root).restore(no_slots_path)
+ with self.assertRaises(AssertionError):
+ no_slot_status.assert_consumed()
+ new_root.var = checkpointable_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.AdamOptimizer(0.1)
+ slot_status.assert_existing_objects_matched()
+ with self.assertRaisesRegexp(AssertionError, "beta1_power"):
+ 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(name="m", var=new_root.var)))
+ else:
+ self.assertIs(new_root.optimizer.get_slot(name="m", var=new_root.var),
+ None)
+ if context.executing_eagerly():
+ new_root.optimizer.minimize(new_root.var.read_value)
+ else:
+ train_op = new_root.optimizer.minimize(new_root.var)
+ # 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()
+ self.assertEqual(14., self.evaluate(
+ new_root.optimizer.get_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 = tracking.Checkpointable()
+ obj.var = variable_scope.get_variable(name="v", initializer=0.)
+ obj.opt = adam.AdamOptimizer(0.1)
+ obj.opt.minimize(obj.var.read_value())
+ self.evaluate(checkpointable_utils.gather_initializers(obj))
+ saver = checkpointable_utils.CheckpointableSaver(obj)
+ saver.save(checkpoint_prefix)
+ before_ops = graph.get_operations()
+ saver.save(checkpoint_prefix)
+ self.assertEqual(before_ops, graph.get_operations())
+
+ 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 = tracking.Checkpointable()
+ obj.var = variable_scope.get_variable(name="v", initializer=0.)
+ obj.opt = adam.AdamOptimizer(0.1)
+ obj.opt.minimize(obj.var.read_value())
+ self.evaluate(checkpointable_utils.gather_initializers(obj))
+ saver = checkpointable_utils.CheckpointableSaver(obj)
+ save_path = saver.save(checkpoint_prefix)
+ saver.restore(save_path)
+ before_ops = graph.get_operations()
+ saver.restore(save_path)
+ self.assertEqual(before_ops, graph.get_operations())
+
+ def testMultipleGraphsNonSlotVariables(self):
+ with context.graph_mode():
+ checkpoint_directory = self.get_temp_dir()
+ checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+ optimizer = adam.AdamOptimizer(0.001)
+ # Construct a model in one graph
+ first_graph = ops.Graph()
+ first_session = session_lib.Session(graph=first_graph)
+ with first_graph.as_default(), first_session.as_default():
+ first_variable = resource_variable_ops.ResourceVariable([1.])
+ first_root_checkpointable = checkpointable_utils.Checkpoint(
+ optimizer=optimizer, variable=first_variable)
+ train_op = optimizer.minimize(first_variable.read_value)
+ self.evaluate(checkpointable_utils.gather_initializers(
+ first_root_checkpointable))
+ self.evaluate(train_op)
+ self.evaluate(first_variable.assign([1.]))
+ self.evaluate(optimizer.get_slot(
+ var=first_variable, name="m").assign([2.]))
+ beta1_power, _ = optimizer._get_beta_accumulators()
+ self.evaluate(beta1_power.assign(3.))
+
+ # Save and load in a second graph
+ second_graph = ops.Graph()
+ with second_graph.as_default(), session_lib.Session(graph=second_graph):
+ second_variable = resource_variable_ops.ResourceVariable([1.])
+ second_root_checkpointable = checkpointable_utils.Checkpoint(
+ optimizer=optimizer, variable=second_variable)
+ train_op = optimizer.minimize(second_variable.read_value)
+ second_root_checkpointable.restore(None).initialize_or_restore()
+ self.evaluate(train_op)
+ self.evaluate(second_variable.assign([4.]))
+ self.evaluate(optimizer.get_slot(
+ var=second_variable, name="m").assign([5.]))
+ beta1_power, _ = optimizer._get_beta_accumulators()
+ self.evaluate(beta1_power.assign(6.))
+ save_path = second_root_checkpointable.save(checkpoint_prefix)
+ self.evaluate(second_variable.assign([7.]))
+ self.evaluate(optimizer.get_slot(
+ var=second_variable, name="m").assign([8.]))
+ beta1_power, _ = optimizer._get_beta_accumulators()
+ self.assertAllEqual(6., self.evaluate(beta1_power))
+ status = second_root_checkpointable.restore(save_path)
+ status.assert_consumed().run_restore_ops()
+ self.assertAllEqual([4.], self.evaluate(second_variable))
+ self.assertAllEqual([5.], self.evaluate(optimizer.get_slot(
+ var=second_variable, name="m")))
+ beta1_power, _ = optimizer._get_beta_accumulators()
+ self.assertAllEqual(6., self.evaluate(beta1_power))
+
+ # Check that the first graph is unmolested
+ with first_graph.as_default(), first_session.as_default():
+ self.assertAllEqual([1.], self.evaluate(first_variable))
+ self.assertAllEqual([2.], self.evaluate(optimizer.get_slot(
+ var=first_variable, name="m")))
+ 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 = checkpointable_utils.Checkpoint(
+ model=model, # Do not save the optimizer with the checkpoint.
+ global_step=training_util.get_or_create_global_step())
+ optimizer_checkpoint = checkpointable_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 = checkpointable_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 = checkpointable_utils.Checkpoint(
+ optimizer=optimizer, model=model,
+ global_step=training_util.get_or_create_global_step())
+ opt_root = checkpointable_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.Checkpointable):
+
+ def __call__(self):
+ with variable_scope.variable_scope("ManualScope") as vs:
+ self.variable_scope = vs
+ with checkpointable_utils.capture_dependencies(template=self):
+ return self._build()
+
+ def _build(self):
+ return variable_scope.get_variable(name="in_manual_scope", shape=[])
+
+
+class TemplateTests(test.TestCase):
+
+ @test_util.run_in_graph_and_eager_modes
+ def test_checkpointable_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,
+ [v1_save, v2_save, manual_scope, manual_scope_v, save_template],
+ checkpointable_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.AdamOptimizer(0.0)
+ save_root = checkpointable_utils.Checkpoint(
+ my_template=save_template, optimizer=optimizer)
+ optimizer.minimize(v1_save.read_value)
+ self.evaluate([v.initializer for v in save_template.variables])
+ 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.AdamOptimizer(0.0)
+ load_root = checkpointable_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)
+ self.assertEqual(3, len(load_template._checkpoint_dependencies))
+ 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.AdamOptimizer(0.001)
+ optimizer_step = training_util.get_or_create_global_step()
+ root_checkpointable = checkpointable_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(checkpointable_utils.gather_initializers(
+ root_checkpointable))
+ 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_checkpointable
+
+ def _set_sentinels(self, root_checkpointable):
+ self.evaluate(root_checkpointable.model._named_dense.bias.assign([101.]))
+ self.evaluate(
+ root_checkpointable.optimizer.get_slot(
+ var=root_checkpointable.model._named_dense.bias, name="m")
+ .assign([102.]))
+ beta1_power, _ = root_checkpointable.optimizer._get_beta_accumulators()
+ self.evaluate(beta1_power.assign(103.))
+
+ def _check_sentinels(self, root_checkpointable):
+ self.assertAllEqual(
+ [1.], self.evaluate(root_checkpointable.model._named_dense.bias))
+ self.assertAllEqual([2.], self.evaluate(
+ root_checkpointable.optimizer.get_slot(
+ var=root_checkpointable.model._named_dense.bias, name="m")))
+ beta1_power, _ = root_checkpointable.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 = checkpointable_utils.CheckpointableSaver(root)
+ self._set_sentinels(root)
+ status = object_saver.restore(save_path)
+ if context.executing_eagerly():
+ self._check_sentinels(root)
+ if context.executing_eagerly():
+ with self.assertRaisesRegexp(AssertionError, "OBJECT_CONFIG_JSON"):
+ status.assert_consumed()
+ with self.assertRaisesRegexp(AssertionError, "OBJECT_CONFIG_JSON"):
+ status.assert_existing_objects_matched()
+ with self.assertRaisesRegexp(AssertionError, "OBJECT_CONFIG_JSON"):
+ 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) as session:
+ root = self._initialized_model()
+ save_path = root.save(session=session, 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/evaluation.py b/tensorflow/python/training/evaluation.py
index a10178f..37d4679 100644
--- a/tensorflow/python/training/evaluation.py
+++ b/tensorflow/python/training/evaluation.py
@@ -253,7 +253,7 @@
if isinstance(h, (_StopAfterNEvalsHook, _MultiStepStopAfterNEvalsHook)):
h._set_evals_completed_tensor(eval_step_value) # pylint: disable=protected-access
- logging.info('Starting evaluation at ' + time.strftime('%Y-%m-%d-%H:%M:%S',
+ logging.info('Starting evaluation at ' + time.strftime('%Y-%m-%dT%H:%M:%SZ',
time.gmtime()))
# Prepare the session creator.
diff --git a/tensorflow/python/training/input_test.py b/tensorflow/python/training/input_test.py
index a3d268a..5efc15d 100644
--- a/tensorflow/python/training/input_test.py
+++ b/tensorflow/python/training/input_test.py
@@ -58,7 +58,7 @@
question = inp.match_filenames_once(
os.path.join(self.get_temp_dir(), "match_filenames.?"))
one = inp.match_filenames_once(additional[1])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
self.assertItemsEqual(
map(compat.as_bytes, filenames), self.evaluate(star))
@@ -84,7 +84,7 @@
with self.cached_session():
love_me = constant_op.constant("Love Me")
love_me_two_times = inp.limit_epochs(love_me, num_epochs=2)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
self.assertEqual(b"Love Me", self.evaluate(love_me_two_times))
self.assertEqual(b"Love Me", self.evaluate(love_me_two_times))
@@ -105,7 +105,7 @@
input_tensor, num_epochs=num_epochs, shuffle=False)
dequeue_many = queue.dequeue_many(len(input_tensor) * num_epochs)
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -132,7 +132,7 @@
input_tensor, element_shape=[4], num_epochs=num_epochs, shuffle=False)
dequeue_many = queue.dequeue_many(len(input_value) * num_epochs)
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -163,7 +163,7 @@
strings, num_epochs=num_epochs, shuffle=False)
dequeue_many = queue.dequeue_many(len(strings) * num_epochs)
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -186,7 +186,7 @@
strings, num_epochs=num_epochs, shuffle=True, seed=271828)
dequeue_many = queue.dequeue_many(len(strings))
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -234,7 +234,7 @@
constant_op.constant(
[], dtype=dtypes.string))
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners(coord=coord)
with self.assertRaises(errors_impl.OutOfRangeError):
@@ -284,7 +284,7 @@
range_size, num_epochs=num_epochs, shuffle=False)
dequeue_many = queue.dequeue_many(range_size * num_epochs)
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -307,7 +307,7 @@
range_size, num_epochs=num_epochs, shuffle=True, seed=314159)
dequeue_many = queue.dequeue_many(range_size)
dequeue = queue.dequeue()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -358,7 +358,7 @@
source_ints = [2, 3, 5, 7]
slices = inp.slice_input_producer(
[source_strings, source_ints], num_epochs=num_epochs, shuffle=False)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -386,7 +386,7 @@
num_epochs=num_epochs,
shuffle=True,
seed=161803)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -487,7 +487,7 @@
batched = inp.batch(
[counter, sparse_counter, "string"], batch_size=batch_size)
batched_fetch = batched
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -555,7 +555,7 @@
counter = examples.count_up_to(num_batches * batch_size)
string = array_ops.tile(["string"],
math_ops.to_int32(array_ops.stack([counter])))
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
batched = inp.batch(
[counter, string], batch_size=batch_size, dynamic_pad=True)
@@ -590,7 +590,7 @@
dense_shape=[1])
pre_batched = inp.batch([counter, sparse_counter, "string"], batch_size=2)
batched = inp.batch(pre_batched, enqueue_many=True, batch_size=batch_size)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -629,7 +629,7 @@
[counter, sparse_counter, "string"],
batch_size=batch_size,
num_threads=4)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -672,7 +672,7 @@
[counter, sparse_counter, "string"],
batch_size=batch_size,
allow_smaller_final_batch=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -730,7 +730,7 @@
batch_size=batch_size,
num_threads=4,
allow_smaller_final_batch=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -872,19 +872,19 @@
for thread in threads:
thread.join()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInput(self):
self._testKeepInputHelper(1, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(1, True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInput(self):
self._testKeepInputHelper(5, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(5, True)
@@ -1058,7 +1058,7 @@
batched_fetch[1].dense_shape.get_shape().as_list())
self.assertAllEqual((batch_size,), batched_fetch[2].get_shape().as_list())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1157,7 +1157,7 @@
self.assertAllEqual((batch_size,), batched[0].get_shape().as_list())
self.assertAllEqual((batch_size, None), batched[1].get_shape().as_list())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1244,7 +1244,7 @@
self.assertAllEqual((2,), batched[1].dense_shape.get_shape().as_list())
self.assertAllEqual((None,), batched[2].get_shape().as_list())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1339,7 +1339,7 @@
self.assertAllEqual((None,), batched[0].get_shape().as_list())
self.assertAllEqual((None, None), batched[1].get_shape().as_list())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1482,19 +1482,19 @@
for thread in threads:
thread.join()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInput(self):
self._testKeepInputHelper(1, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(1, True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInput(self):
self._testKeepInputHelper(5, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(5, True)
@@ -1644,7 +1644,7 @@
min_after_dequeue=16,
seed=141421)
batched_fetch = batched
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1702,7 +1702,7 @@
seed=141421,
allow_smaller_final_batch=True)
batched_fetch = batched
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1756,7 +1756,7 @@
min_after_dequeue=16,
seed=173205,
num_threads=4)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1807,7 +1807,7 @@
seed=173205,
num_threads=4,
allow_smaller_final_batch=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -1905,19 +1905,19 @@
for thread in threads:
thread.join()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInput(self):
self._testKeepInputHelper(1, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(1, True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInput(self):
self._testKeepInputHelper(5, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(5, True)
@@ -2070,7 +2070,7 @@
batched_fetch[1].dense_shape.get_shape().as_list())
self.assertAllEqual((batch_size,), batched_fetch[2].get_shape().as_list())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -2165,7 +2165,7 @@
self.assertAllEqual((2,), batched[1].dense_shape.get_shape().as_list())
self.assertAllEqual((None,), batched[2].get_shape().as_list())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
variables.local_variables_initializer().run()
threads = queue_runner_impl.start_queue_runners()
@@ -2309,19 +2309,19 @@
for thread in threads:
thread.join()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInput(self):
self._testKeepInputHelper(1, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSingleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(1, True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInput(self):
self._testKeepInputHelper(5, False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultipleThreadKeepInputEnqueueMany(self):
self._testKeepInputHelper(5, True)
diff --git a/tensorflow/python/training/learning_rate_decay_test.py b/tensorflow/python/training/learning_rate_decay_test.py
index 9de5bc8..1029d4c 100644
--- a/tensorflow/python/training/learning_rate_decay_test.py
+++ b/tensorflow/python/training/learning_rate_decay_test.py
@@ -101,7 +101,7 @@
self.assertAllClose(self.evaluate(decayed_lr), 0.001, 1e-6)
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPiecewiseConstantEdgeCases(self):
x_int = resource_variable_ops.ResourceVariable(
0, dtype=variables.dtypes.int32)
diff --git a/tensorflow/python/training/localhost_cluster_performance_test.py b/tensorflow/python/training/localhost_cluster_performance_test.py
index 7c097b9..c4cbc8a 100644
--- a/tensorflow/python/training/localhost_cluster_performance_test.py
+++ b/tensorflow/python/training/localhost_cluster_performance_test.py
@@ -25,6 +25,7 @@
from tensorflow.python.client import session as session_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import partitioned_variables
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
@@ -34,6 +35,7 @@
class CreateLocalClusterTest(test.TestCase):
+ @test_util.run_v1_only("b/120545219")
def testCreateLocalCluster(self):
workers, _ = test.create_local_cluster(num_workers=2, num_ps=2)
worker_sessions = [session_lib.Session(w.target) for w in workers]
diff --git a/tensorflow/python/training/monitored_session_test.py b/tensorflow/python/training/monitored_session_test.py
index 9dbcfa5..99ee9ea 100644
--- a/tensorflow/python/training/monitored_session_test.py
+++ b/tensorflow/python/training/monitored_session_test.py
@@ -541,6 +541,7 @@
self.assertFalse(wrapped_sess1.should_stop())
self.assertTrue(wrapped_sess1.should_stop())
+ @test_util.run_deprecated_v1
def test_close_twice(self):
with self.cached_session() as sess:
wrapped_sess = monitored_session._WrappedSession(sess)
diff --git a/tensorflow/python/training/moving_averages.py b/tensorflow/python/training/moving_averages.py
index 8785f9a..72670f0 100644
--- a/tensorflow/python/training/moving_averages.py
+++ b/tensorflow/python/training/moving_averages.py
@@ -17,6 +17,7 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.distribute import reduce_util as ds_reduce_util
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
@@ -26,7 +27,6 @@
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
-from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.training import slot_creator
from tensorflow.python.util.tf_export import tf_export
diff --git a/tensorflow/python/training/moving_averages_test.py b/tensorflow/python/training/moving_averages_test.py
index b15f737..0a7cff4 100644
--- a/tensorflow/python/training/moving_averages_test.py
+++ b/tensorflow/python/training/moving_averages_test.py
@@ -43,7 +43,7 @@
decay = 0.25
assign = moving_averages.assign_moving_average(
var, val, decay, zero_debias=False)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose([10.0, 11.0], self.evaluate(var))
assign.op.run()
self.assertAllClose(
@@ -57,7 +57,7 @@
val = constant_op.constant([1.0, 2.0], dtypes.float32)
decay = 0.25
assign = moving_averages.assign_moving_average(var, val, decay)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllClose([0.0, 0.0], self.evaluate(var))
assign.op.run()
self.assertAllClose(
@@ -98,7 +98,7 @@
val = array_ops.placeholder(dtypes.float32, [])
wma = moving_averages.weighted_moving_average(val, decay, weight)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Get the first weighted moving average.
val_1 = 3.0
@@ -125,7 +125,7 @@
val = array_ops.placeholder(dtypes.bfloat16, [])
wma = moving_averages.weighted_moving_average(val, decay, weight)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Get the first weighted moving average.
val_1 = 3.0
@@ -164,7 +164,7 @@
thirties = _Repeat(30.0, dim)
var0 = variables.Variable(tens, name="v0")
var1 = variables.Variable(thirties, name="v1")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Note that tensor2 is not a Variable but just a plain Tensor resulting
# from the sum operation.
tensor2 = var0 + var1
@@ -178,7 +178,7 @@
self.assertFalse(avg0 in variables.trainable_variables())
self.assertFalse(avg1 in variables.trainable_variables())
self.assertFalse(avg2 in variables.trainable_variables())
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("v0/ExponentialMovingAverage:0", avg0.name)
self.assertEqual("v1/ExponentialMovingAverage:0", avg1.name)
@@ -219,38 +219,38 @@
(10.0 + 30.0) * (1 - dk)) / _Scale(dk, 2), dim)
self.assertAllClose(expected, self.evaluate(avg2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNoNumUpdates_Scalar(self):
with self.cached_session():
ema = moving_averages.ExponentialMovingAverage(0.25)
self._CheckDecay(ema, actual_decay=0.25, dim=1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNoNumUpdates_Scalar_Debias(self):
with self.cached_session():
ema = moving_averages.ExponentialMovingAverage(0.25, zero_debias=True)
self._CheckDecay(ema, actual_decay=0.25, dim=1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNoNumUpdates_Vector(self):
with self.cached_session():
ema = moving_averages.ExponentialMovingAverage(0.25)
self._CheckDecay(ema, actual_decay=0.25, dim=5)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNoNumUpdates_Vector_Debias(self):
with self.cached_session():
ema = moving_averages.ExponentialMovingAverage(0.25, zero_debias=True)
self._CheckDecay(ema, actual_decay=0.25, dim=5)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNumUpdates_Scalar(self):
with self.cached_session():
# With num_updates 1, the decay applied is 0.1818
ema = moving_averages.ExponentialMovingAverage(0.25, num_updates=1)
self._CheckDecay(ema, actual_decay=0.181818, dim=1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNumUpdates_Scalar_Debias(self):
with self.cached_session():
# With num_updates 1, the decay applied is 0.1818
@@ -258,14 +258,14 @@
0.25, num_updates=1, zero_debias=True)
self._CheckDecay(ema, actual_decay=0.181818, dim=1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNumUpdates_Vector(self):
with self.cached_session():
# With num_updates 1, the decay applied is 0.1818
ema = moving_averages.ExponentialMovingAverage(0.25, num_updates=1)
self._CheckDecay(ema, actual_decay=0.181818, dim=5)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNumUpdates_Vector_Debias(self):
with self.cached_session():
# With num_updates 1, the decay applied is 0.1818
@@ -273,7 +273,7 @@
0.25, num_updates=1, zero_debias=True)
self._CheckDecay(ema, actual_decay=0.181818, dim=5)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesWithControlDeps(self):
with self.cached_session() as sess:
v0 = variables.Variable(0, name="v0")
@@ -299,7 +299,7 @@
self.assertEqual([17.5], self.evaluate(v1_avg))
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBasicEager(self):
v0 = variables.Variable(1.0)
v1 = variables.Variable(2.0)
@@ -355,11 +355,11 @@
self.assertEqual(ema.average(v1).op.name, ema.average_name(v1))
self.assertEqual(ema.average(tensor2).op.name, ema.average_name(tensor2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNames(self):
self.averageVariablesNamesHelper(zero_debias=True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNamesNoDebias(self):
self.averageVariablesNamesHelper(zero_debias=False)
@@ -405,15 +405,15 @@
self.assertEqual(
ema.average(tensor2).op.name, ema.average_name(tensor2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNamesRespectScope(self):
self.averageVariablesNamesRespectScopeHelper(zero_debias=True)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesNamesRespectScopeNoDebias(self):
self.averageVariablesNamesRespectScopeHelper(zero_debias=False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSubsetAverageVariablesNames(self):
with self.cached_session():
v0 = variables.Variable(10.0, name="v0")
@@ -442,7 +442,7 @@
self.assertEqual(ema.average(v1).op.name, ema.average_name(v1))
self.assertEqual(ema.average(tensor2).op.name, ema.average_name(tensor2))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAverageVariablesDeviceAssignment(self):
with ops.device("/job:dev_v0"):
v0 = variables.Variable(10.0, name="v0")
diff --git a/tensorflow/python/training/optimizer.py b/tensorflow/python/training/optimizer.py
index a9508b8..eaa563e 100644
--- a/tensorflow/python/training/optimizer.py
+++ b/tensorflow/python/training/optimizer.py
@@ -25,6 +25,7 @@
import six
from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribution_strategy_context as distribute_ctx
from tensorflow.python.distribute import reduce_util as ds_reduce_util
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
@@ -38,7 +39,6 @@
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
-from tensorflow.python.training import distribution_strategy_context as distribute_ctx
from tensorflow.python.training import slot_creator
from tensorflow.python.training.checkpointable import base as checkpointable
from tensorflow.python.util import nest
@@ -822,7 +822,10 @@
name=name, shape=None)
if restored_initial_value is not None:
initial_value = restored_initial_value
- v = variable_scope.variable(initial_value, name=name, trainable=False)
+ v = variable_scope.variable(
+ initial_value, name=name, trainable=False,
+ use_resource=resource_variable_ops.is_resource_variable(
+ colocate_with))
# Restore this variable by name if necessary, but don't add a
# Checkpointable dependency. Optimizers return the current graph's
# non-slot variables from _checkpoint_dependencies explicitly rather
diff --git a/tensorflow/python/training/quantize_training_test.py b/tensorflow/python/training/quantize_training_test.py
index 62e783f..2352af7 100644
--- a/tensorflow/python/training/quantize_training_test.py
+++ b/tensorflow/python/training/quantize_training_test.py
@@ -53,7 +53,7 @@
# Test that save/restoring works for EMA variables generated in the
# quantized training rewrite.
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testQuantizedSaveRestore(self):
save_path = os.path.join(self.get_temp_dir(), 'quantized_save_restore')
diff --git a/tensorflow/python/training/queue_runner_test.py b/tensorflow/python/training/queue_runner_test.py
index 4113cec..c508507 100644
--- a/tensorflow/python/training/queue_runner_test.py
+++ b/tensorflow/python/training/queue_runner_test.py
@@ -41,7 +41,7 @@
class QueueRunnerTest(test.TestCase):
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBasic(self):
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
@@ -49,7 +49,7 @@
var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
threads = qr.create_threads(sess)
self.assertEqual(sorted(t.name for t in threads),
@@ -62,7 +62,7 @@
# The variable should be 3.
self.assertEqual(3, self.evaluate(var))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testTwoOps(self):
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
@@ -77,7 +77,7 @@
self.assertEqual(sorted(t.name for t in threads),
["QueueRunnerThread-fifo_queue-CountUpTo:0",
"QueueRunnerThread-fifo_queue-CountUpTo_1:0"])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
for t in threads:
t.start()
for t in threads:
@@ -93,7 +93,7 @@
qr = queue_runner_impl.QueueRunner(queue, [_MockOp("i fail"),
_MockOp("so fail")])
threads = qr.create_threads(sess)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
for t in threads:
t.start()
for t in threads:
@@ -132,7 +132,7 @@
with self.assertRaisesRegexp(errors_impl.OutOfRangeError, "is closed"):
self.evaluate(dequeue1)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRespectCoordShouldStop(self):
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
@@ -140,7 +140,7 @@
var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
# As the coordinator to stop. The queue runner should
# finish immediately.
@@ -196,7 +196,7 @@
var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
coord = coordinator.Coordinator()
qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
# NOTE that this test does not actually start the threads.
@@ -212,7 +212,7 @@
var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
coord = coordinator.Coordinator()
qr = queue_runner_impl.QueueRunner(queue, [count_up_to])
threads = []
@@ -221,7 +221,7 @@
new_threads = qr.create_threads(sess, coord=coord)
self.assertEqual([], new_threads)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testThreads(self):
with self.cached_session() as sess:
# CountUpTo will raise OUT_OF_RANGE when it reaches the count.
@@ -229,7 +229,7 @@
var = variables.VariableV1(zero64)
count_up_to = var.count_up_to(3)
queue = data_flow_ops.FIFOQueue(10, dtypes.float32)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
qr = queue_runner_impl.QueueRunner(queue, [count_up_to,
_MockOp("bad_op")])
threads = qr.create_threads(sess, start=True)
diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py
index 4cd09f8..348b8bf 100644
--- a/tensorflow/python/training/saver.py
+++ b/tensorflow/python/training/saver.py
@@ -14,7 +14,11 @@
# ==============================================================================
# pylint: disable=invalid-name
-"""Save and restore variables."""
+"""Save and restore variables.
+
+Symbols in this file are deprecated. See replacements in
+tensorflow/python/training/checkpointable and tensorflow/python/training/saving.
+"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
@@ -25,7 +29,6 @@
import uuid
import numpy as np
-import six
from tensorflow.core.protobuf import checkpointable_object_graph_pb2
from tensorflow.core.protobuf import meta_graph_pb2
@@ -42,16 +45,15 @@
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_io_ops
from tensorflow.python.ops import io_ops
-from tensorflow.python.ops import resource_variable_ops
-from tensorflow.python.ops import state_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import checkpoint_management
-from tensorflow.python.training import saveable_object
from tensorflow.python.training import training_util
from tensorflow.python.training.checkpointable import base as checkpointable
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.util import compat
from tensorflow.python.util.tf_export import tf_export
@@ -67,31 +69,6 @@
remove_checkpoint = checkpoint_management.remove_checkpoint
-# Op names which identify variable reads which should be saved.
-_VARIABLE_OPS = set(["Variable",
- "VariableV2",
- "AutoReloadVariable",
- "VarHandleOp",
- "ReadVariableOp"])
-
-
-def _set_cpu0(device_string):
- """Creates a new device string based on `device_string` but using /CPU:0.
-
- If the device is already on /CPU:0, this is a no-op.
-
- Args:
- device_string: A device string.
-
- Returns:
- A device string.
- """
- parsed_device = pydev.DeviceSpec.from_string(device_string)
- parsed_device.device_type = "CPU"
- parsed_device.device_index = 0
- return parsed_device.to_string()
-
-
class BaseSaverBuilder(object):
"""Base class for Savers.
@@ -101,64 +78,9 @@
SaveSpec = saveable_object.SaveSpec
SaveableObject = saveable_object.SaveableObject
- class VariableSaveable(SaveableObject):
- """SaveableObject implementation that handles Variables."""
-
- def __init__(self, var, slice_spec, name):
- spec = BaseSaverBuilder.SaveSpec(var, slice_spec, name, dtype=var.dtype)
- super(BaseSaverBuilder.VariableSaveable, self).__init__(var, [spec], name)
-
- def restore(self, restored_tensors, restored_shapes):
- restored_tensor = restored_tensors[0]
- if restored_shapes is not None:
- restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
- return state_ops.assign(
- self.op,
- restored_tensor,
- validate_shape=restored_shapes is None and
- self.op.get_shape().is_fully_defined())
-
- class ResourceVariableSaveable(SaveableObject):
- """SaveableObject implementation that handles ResourceVariables."""
-
- def __init__(self, var, slice_spec, name):
- self._var_device = var.device
- self._var_shape = var.shape
- if isinstance(var, ops.Tensor):
- self.handle_op = var.op.inputs[0]
- tensor = var
- elif isinstance(var, resource_variable_ops.ResourceVariable):
-
- def _read_variable_closure(v):
- def f():
- with ops.device(v.device):
- x = v.read_value()
- # To allow variables placed on non-CPU devices to be checkpointed,
- # we copy them to CPU on the same machine first.
- with ops.device("/device:CPU:0"):
- return array_ops.identity(x)
- return f
-
- self.handle_op = var.handle
- tensor = _read_variable_closure(var)
- else:
- raise ValueError(
- "Saveable is neither a resource variable nor a read operation."
- " Got: %s" % repr(var))
- spec = BaseSaverBuilder.SaveSpec(tensor, slice_spec, name,
- dtype=var.dtype)
- super(BaseSaverBuilder.ResourceVariableSaveable, self).__init__(
- var, [spec], name)
-
- def restore(self, restored_tensors, restored_shapes):
- restored_tensor = restored_tensors[0]
- if restored_shapes is not None:
- restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
- # Copy the restored tensor to the variable's device.
- with ops.device(self._var_device):
- restored_tensor = array_ops.identity(restored_tensor)
- return resource_variable_ops.shape_safe_assign_variable_handle(
- self.handle_op, self._var_shape, restored_tensor)
+ # Aliases for code which was moved but still has lots of users.
+ VariableSaveable = saveable_object_util.ReferenceVariableSaveable
+ ResourceVariableSaveable = saveable_object_util.ResourceVariableSaveable
def __init__(self, write_version=saver_pb2.SaverDef.V2):
self._write_version = write_version
@@ -224,7 +146,11 @@
del restore_sequentially
all_tensors = []
for saveable in saveables:
- with ops.device(_set_cpu0(saveable.device) if saveable.device else None):
+ if saveable.device:
+ device = saveable_object_util.set_cpu0(saveable.device)
+ else:
+ device = None
+ with ops.device(device):
all_tensors.extend(
self.restore_op(filename_tensor, saveable, preferred_shard))
return all_tensors
@@ -336,7 +262,7 @@
last_device = None
for shard, (device, saveables) in enumerate(per_device):
last_device = device
- with ops.device(_set_cpu0(device)):
+ with ops.device(saveable_object_util.set_cpu0(device)):
sharded_filename = self.sharded_filename(tmp_checkpoint_prefix, shard,
num_shards_tensor)
sharded_prefixes.append(sharded_filename)
@@ -344,7 +270,7 @@
with ops.control_dependencies([x.op for x in sharded_saves]):
# Co-locates the merge step with the last device.
- with ops.device(_set_cpu0(last_device)):
+ with ops.device(saveable_object_util.set_cpu0(last_device)):
# V2 format write path consists of a metadata merge step. Once merged,
# attempts to delete the temporary directory, "<user-fed prefix>_temp".
merge_step = gen_io_ops.merge_v2_checkpoints(
@@ -459,10 +385,6 @@
name="restore_shard"))
return control_flow_ops.group(*sharded_restores, name="restore_all")
- @staticmethod
- def _IsVariable(v):
- return isinstance(v, ops.Tensor) and v.op.type in _VARIABLE_OPS
-
def _GroupByDevices(self, saveables):
"""Group Variable tensor slices per device.
@@ -490,220 +412,6 @@
per_device[canonical_device.pop()].append(saveable)
return sorted(per_device.items(), key=lambda t: t[0])
- @staticmethod
- def OpListToDict(op_list, convert_variable_to_tensor=True):
- """Create a dictionary of names to operation lists.
-
- Args:
- op_list: A list, tuple, or set of Variables or SaveableObjects.
- convert_variable_to_tensor: Whether or not to convert single Variables
- with no slice info into Tensors.
-
- Returns:
- A dictionary of names to the operations that must be saved under
- that name. Variables with save_slice_info are grouped together under the
- same key in no particular order.
-
- Raises:
- TypeError: If the type of op_list or its elements is not supported.
- ValueError: If at least two saveables share the same name.
- """
- if not isinstance(op_list, (list, tuple, set)):
- raise TypeError("Variables to save should be passed in a dict or a "
- "list: %s" % op_list)
- # When ResourceVariables are converted to Tensors, read ops are added to the
- # graph. Sorting the op_list ensures that the resulting graph is always
- # constructed in a deterministic way:
- op_list = sorted(op_list, key=lambda x: x.name)
- names_to_saveables = {}
- # pylint: disable=protected-access
- for var in op_list:
- if isinstance(var, BaseSaverBuilder.SaveableObject):
- names_to_saveables[var.name] = var
- elif isinstance(var, variables.PartitionedVariable):
- if var.name in names_to_saveables:
- raise ValueError("At least two variables have the same name: %s" %
- var.name)
- names_to_saveables[var.name] = var
- elif isinstance(var, variables.Variable) and var._save_slice_info:
- name = var._save_slice_info.full_name
- if name in names_to_saveables:
- if not isinstance(names_to_saveables[name], list):
- raise ValueError("Mixing slices and non-slices with the same name: "
- "%s" % name)
- names_to_saveables[name].append(var)
- else:
- names_to_saveables[name] = [var]
- elif (isinstance(var, checkpointable.CheckpointableBase)
- and not isinstance(var, variables.Variable)):
- checkpointable_saveables = [
- (factory() if callable(factory) else factory)
- for factory in var._gather_saveables_for_checkpoint().values()]
- names_to_saveables.update(
- BaseSaverBuilder.OpListToDict(checkpointable_saveables))
- else:
- if context.executing_eagerly():
- if not isinstance(var, resource_variable_ops.ResourceVariable):
- raise ValueError(
- "Can only save/restore ResourceVariables when eager execution "
- "is enabled, type: %s." % type(var))
- set_var = names_to_saveables.setdefault(var._shared_name, var)
- if set_var is not var:
- raise ValueError(
- ("Two different ResourceVariable objects with the same "
- "shared_name '%s' were passed to the Saver. This likely means "
- "that they were created in different Graphs or isolation "
- "contexts, and may not be checkpointed together.") %
- (var._shared_name,))
- else:
- if convert_variable_to_tensor:
- if isinstance(var, resource_variable_ops.ResourceVariable):
- var = var._graph_element # pylint: disable=protected-access
- else:
- var = ops.internal_convert_to_tensor(var, as_ref=True)
- if not BaseSaverBuilder._IsVariable(var):
- raise TypeError("Variable to save is not a Variable: %s" % var)
- if var.op.type == "ReadVariableOp":
- name = var.op.inputs[0].op.name
- else:
- name = var.op.name
- if name in names_to_saveables:
- raise ValueError("At least two variables have the same name: %s" %
- name)
- names_to_saveables[name] = var
-
- # pylint: enable=protected-access
- return names_to_saveables
-
- @staticmethod
- def SaveableObjectsForOp(op, name):
- """Create `SaveableObject`s from an operation.
-
- Args:
- op: A variable, operation, or SaveableObject to coerce into a
- SaveableObject.
- name: A string name for the SaveableObject.
-
- Yields:
- `SaveableObject`s which together save/restore `op`.
-
- Raises:
- TypeError: If `name` is not a string.
- ValueError: For operations with no known conversion to SaveableObject.
- """
- if not isinstance(name, six.string_types):
- raise TypeError(
- "names_to_saveables must be a dict mapping string names to "
- "checkpointable operations. Name is not a string: %s" % name)
- if isinstance(op, BaseSaverBuilder.SaveableObject):
- yield op
- elif isinstance(op, (list, tuple, variables.PartitionedVariable)):
- if isinstance(op, variables.PartitionedVariable):
- op = list(op)
- # A set of slices.
- slice_name = None
- # pylint: disable=protected-access
- for variable in op:
- if not isinstance(variable, variables.Variable):
- raise ValueError("Slices must all be Variables: %s" % variable)
- if not variable._save_slice_info:
- raise ValueError("Slices must all be slices: %s" % variable)
- if slice_name is None:
- slice_name = variable._save_slice_info.full_name
- elif slice_name != variable._save_slice_info.full_name:
- raise ValueError(
- "Slices must all be from the same tensor: %s != %s" %
- (slice_name, variable._save_slice_info.full_name))
- if variable.op.type in ["Variable", "VariableV2",
- "AutoReloadVariable"]:
- yield BaseSaverBuilder.VariableSaveable(
- variable, variable._save_slice_info.spec, name)
- else:
- yield BaseSaverBuilder.ResourceVariableSaveable(
- variable, variable._save_slice_info.spec, name)
- # pylint: enable=protected-access
- elif isinstance(op, checkpointable.CheckpointableBase) and not isinstance(
- op, variables.Variable):
- # pylint: disable=protected-access
- for attr, factory in op._gather_saveables_for_checkpoint().items():
- if attr == checkpointable.VARIABLE_VALUE_KEY:
- # Keep original name for classes masquerading as variables.
- full_name = name
- else:
- full_name = name + "_" + attr
- op = (factory(full_name) if callable(factory) else factory)
- for op in BaseSaverBuilder.SaveableObjectsForOp(op, op.name):
- yield op
- # pylint: enable=protected-access
- else:
- # A variable or tensor.
- if context.executing_eagerly():
- if not isinstance(op, resource_variable_ops.ResourceVariable):
- raise ValueError("Can only save/restore ResourceVariable eager "
- "mode is enabled, type: %s." % type(op))
- yield BaseSaverBuilder.ResourceVariableSaveable(op, "", name)
- else:
- if isinstance(op, resource_variable_ops.ResourceVariable):
- variable = op._graph_element # pylint: disable=protected-access
- else:
- variable = ops.internal_convert_to_tensor(op, as_ref=True)
- if not BaseSaverBuilder._IsVariable(variable):
- raise TypeError("names_to_saveables must be a dict mapping string "
- "names to Tensors/Variables. Not a variable: %s" %
- variable)
- if variable.op.type in ["Variable", "VariableV2",
- "AutoReloadVariable"]:
- yield BaseSaverBuilder.VariableSaveable(variable, "", name)
- else:
- yield BaseSaverBuilder.ResourceVariableSaveable(
- variable, "", name)
-
- def _ValidateAndSliceInputs(self, names_to_saveables):
- """Returns the variables and names that will be used for a Saver.
-
- Args:
- names_to_saveables: A dict (k, v) where k is the name of an operation and
- v is an operation to save or a BaseSaverBuilder.Saver.
-
- Returns:
- A list of BaseSaverBuilder.SaveableObject objects.
-
- Raises:
- TypeError: If any of the keys are not strings or any of the
- values are not one of Tensor or Variable or a checkpointable operation.
- ValueError: If the same operation is given in more than one value
- (this also applies to slices of SlicedVariables).
- """
- if not isinstance(names_to_saveables, dict):
- names_to_saveables = BaseSaverBuilder.OpListToDict(names_to_saveables)
-
- saveables = []
- seen_ops = set()
- for name, op in sorted(names_to_saveables.items(),
- # Avoid comparing ops, sort only by name.
- key=lambda x: x[0]):
- for converted_saveable_object in self.SaveableObjectsForOp(op, name):
- self._AddSaveable(saveables, seen_ops, converted_saveable_object)
- return saveables
-
- def _AddSaveable(self, saveables, seen_ops, saveable):
- """Adds the saveable to the saveables list.
-
- Args:
- saveables: List to append the SaveableObject to.
- seen_ops: Set of the ops of the saveables already processed. Used to
- check that each saveable is only saved once.
- saveable: The saveable.
-
- Raises:
- ValueError: If the saveable has already been processed.
- """
- if saveable.op in seen_ops:
- raise ValueError("The same saveable will be restored with two names: %s" %
- saveable.name)
- saveables.append(saveable)
- seen_ops.add(saveable.op)
-
def build(self,
names_to_saveables,
reshape=False,
@@ -775,7 +483,8 @@
raise ValueError("save and restore operations need to be built together "
" when eager execution is not enabled.")
- saveables = self._ValidateAndSliceInputs(names_to_saveables)
+ saveables = saveable_object_util.validate_and_slice_inputs(
+ names_to_saveables)
if max_to_keep is None:
max_to_keep = 0
@@ -1668,6 +1377,37 @@
NOTE: Restarting training from saved `meta_graph` only works if the
device assignments have not changed.
+ Example 2:
+ Variables, placeholders, and independent operations can also be stored, as
+ shown in the following example.
+
+ ```Python
+ # Saving contents and operations.
+ v1 = tf.placeholder(tf.float32, name="v1")
+ v2 = tf.placeholder(tf.float32, name="v2")
+ v3 = tf.mul(v1, v2)
+ vx = tf.Variable(10.0, name="vx")
+ v4 = tf.add(v3, vx, name="v4")
+ saver = tf.train.Saver([vx])
+ sess = tf.Session()
+ sess.run(tf.initialize_all_variables())
+ sess.run(vx.assign(tf.add(vx, vx)))
+ result = sess.run(v4, feed_dict={v1:12.0, v2:3.3})
+ print(result)
+ saver.save(sess, "./model_ex1")
+ ```
+
+ Later this model can be restored and contents loaded.
+
+ ```Python
+ # Restoring variables and running operations.
+ saver = tf.train.import_meta_graph("./model_ex1.meta")
+ sess = tf.Session()
+ saver.restore(sess, "./model_ex1")
+ result = sess.run("v4:0", feed_dict={"v1:0": 12.0, "v2:0": 3.3})
+ print(result)
+ ```
+
Args:
meta_graph_or_file: `MetaGraphDef` protocol buffer or filename (including
the path) containing a `MetaGraphDef`.
@@ -1910,7 +1650,7 @@
if builder is None:
builder = BulkSaverBuilder()
- saveables = builder._ValidateAndSliceInputs(var_list) # pylint: disable=protected-access
+ saveables = saveable_object_util.validate_and_slice_inputs(var_list)
current_names = set()
for saveable in saveables:
for spec in saveable.specs:
diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py
index 5d621ba..d1b51ad 100644
--- a/tensorflow/python/training/saver_test.py
+++ b/tensorflow/python/training/saver_test.py
@@ -124,8 +124,8 @@
if not context.executing_eagerly():
self.assertEqual(
len(variables.report_uninitialized_variables().eval()), 2)
- self.assertEqual(0, len(v2.keys().eval()))
- self.assertEqual(0, len(v2.values().eval()))
+ self.assertEqual(0, len(self.evaluate(v2.keys())))
+ self.assertEqual(0, len(self.evaluate(v2.values())))
# Restore the saved values in the parameter nodes.
save = saver_module.Saver({"v0": v0, "v1": v1, "v2": v2.saveable})
save.restore(sess, save_path)
@@ -331,10 +331,10 @@
self.evaluate(init_all_op)
# Check that the parameter nodes have been initialized.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
- self.assertEqual(b"k1", v2.keys().eval())
- self.assertEqual(30.0, v2.values().eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
+ self.assertEqual(b"k1", self.evaluate(v2.keys()))
+ self.assertEqual(30.0, self.evaluate(v2.values()))
# Save the initialized values in the file at "save_path"
val = save.save(sess, save_path1)
@@ -360,16 +360,16 @@
# Assert that the variables are not initialized.
self.assertEqual(
len(variables.report_uninitialized_variables().eval()), 2)
- self.assertEqual(0, len(v2.keys().eval()))
- self.assertEqual(0, len(v2.values().eval()))
+ self.assertEqual(0, len(self.evaluate(v2.keys())))
+ self.assertEqual(0, len(self.evaluate(v2.values())))
# Restore the saved values in the parameter nodes.
save.restore(sess, save_path2)
# Check that the parameter nodes have been restored.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
- self.assertEqual(b"k1", v2.keys().eval())
- self.assertEqual(30.0, v2.values().eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
+ self.assertEqual(b"k1", self.evaluate(v2.keys()))
+ self.assertEqual(30.0, self.evaluate(v2.values()))
@test_util.run_deprecated_v1
def testFilenameTensor(self):
@@ -390,7 +390,7 @@
ValueError, "The passed save_path is not a valid checkpoint:"):
save.restore(sess, "invalid path")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInt64(self):
save_path = os.path.join(self.get_temp_dir(), "int64")
@@ -398,7 +398,7 @@
# Build a graph with 1 node, and save and restore for them.
v = variables.VariableV1(np.int64(15), name="v")
save = saver_module.Saver({"v": v}, restore_sequentially=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Save the initialized values in the file at "save_path"
val = save.save(sess, save_path)
@@ -416,7 +416,7 @@
# Restore the saved values in the parameter nodes.
save.restore(sess, save_path)
# Check that the parameter nodes have been restored.
- self.assertEqual(np.int64(15), v.eval())
+ self.assertEqual(np.int64(15), self.evaluate(v))
def testSomeErrors(self):
with ops_lib.Graph().as_default():
@@ -466,7 +466,7 @@
# Verify non-duplicate names work.
saver_module.Saver({"v0": v0, "v2": v2.saveable})
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testBasicsWithListOfVariables(self):
save_path = os.path.join(self.get_temp_dir(), "basics_with_list")
@@ -478,14 +478,14 @@
v2 = saver_test_utils.CheckpointedOp(name="v2")
v2_init = v2.insert("k1", 30.0)
save = saver_module.Saver([v0, v1, v2.saveable])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
v2_init.run()
# Check that the parameter nodes have been initialized.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
- self.assertEqual(b"k1", v2.keys().eval())
- self.assertEqual(30.0, v2.values().eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
+ self.assertEqual(b"k1", self.evaluate(v2.keys()))
+ self.assertEqual(30.0, self.evaluate(v2.values()))
# Save the initialized values in the file at "save_path"
val = save.save(sess, save_path)
@@ -506,16 +506,16 @@
with self.assertRaisesWithPredicateMatch(
errors_impl.OpError, lambda e: "uninitialized value v1" in e.message):
self.evaluate(v1)
- self.assertEqual(0, len(v2.keys().eval()))
- self.assertEqual(0, len(v2.values().eval()))
+ self.assertEqual(0, len(self.evaluate(v2.keys())))
+ self.assertEqual(0, len(self.evaluate(v2.values())))
# Restore the saved values in the parameter nodes.
save.restore(sess, save_path)
# Check that the parameter nodes have been restored.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
- self.assertEqual(b"k1", v2.keys().eval())
- self.assertEqual(30.0, v2.values().eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
+ self.assertEqual(b"k1", self.evaluate(v2.keys()))
+ self.assertEqual(30.0, self.evaluate(v2.values()))
# Build another graph with 2 nodes, initialized
# differently, and a Restore node for them.
@@ -525,20 +525,20 @@
v2_2 = saver_test_utils.CheckpointedOp(name="v2")
save2 = saver_module.Saver([v0_2, v1_2, v2_2.saveable])
v2_2.insert("k1000", 3000.0).run()
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Check that the parameter nodes have been initialized.
- self.assertEqual(1000.0, v0_2.eval())
- self.assertEqual(2000.0, v1_2.eval())
- self.assertEqual(b"k1000", v2_2.keys().eval())
- self.assertEqual(3000.0, v2_2.values().eval())
+ self.assertEqual(1000.0, self.evaluate(v0_2))
+ self.assertEqual(2000.0, self.evaluate(v1_2))
+ self.assertEqual(b"k1000", self.evaluate(v2_2.keys()))
+ self.assertEqual(3000.0, self.evaluate(v2_2.values()))
# Restore the values saved earlier in the parameter nodes.
save2.restore(sess, save_path)
# Check that the parameter nodes have been restored.
- self.assertEqual(10.0, v0_2.eval())
- self.assertEqual(20.0, v1_2.eval())
- self.assertEqual(b"k1", v2_2.keys().eval())
- self.assertEqual(30.0, v2_2.values().eval())
+ self.assertEqual(10.0, self.evaluate(v0_2))
+ self.assertEqual(20.0, self.evaluate(v1_2))
+ self.assertEqual(b"k1", self.evaluate(v2_2.keys()))
+ self.assertEqual(30.0, self.evaluate(v2_2.values()))
def _SaveAndLoad(self, var_name, var_value, other_value, save_path):
with self.session(graph=ops_lib.Graph()) as sess:
@@ -582,14 +582,14 @@
with sess.graph.device(test.gpu_device_name()):
v0_1 = variables.VariableV1(123.45)
save = saver_module.Saver({"v0": v0_1})
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
save.save(sess, save_path)
with session.Session("", graph=ops_lib.Graph()) as sess:
with sess.graph.device(test.gpu_device_name()):
v0_2 = variables.VariableV1(543.21)
save = saver_module.Saver({"v0": v0_2})
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
def testSharedServerOnGPU(self):
if not test.is_gpu_available():
@@ -599,14 +599,14 @@
with sess.graph.device(test.gpu_device_name()):
v0_1 = variables.VariableV1(123.45)
save = saver_module.Saver({"v0": v0_1}, sharded=True, allow_empty=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
save.save(sess, save_path)
with session.Session("", graph=ops_lib.Graph()) as sess:
with sess.graph.device(test.gpu_device_name()):
v0_2 = variables.VariableV1(543.21)
save = saver_module.Saver({"v0": v0_2}, sharded=True, allow_empty=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
def testVariables(self):
save_path = os.path.join(self.get_temp_dir(), "variables")
@@ -627,10 +627,10 @@
# Saver with no arg, defaults to 'all variables'.
save = saver_module.Saver()
save.restore(sess, save_path)
- self.assertAllClose(1.0, one.eval())
- self.assertAllClose([2.0, 2.0, 2.0], twos.eval())
- self.assertEqual(b"k1", v2.keys().eval())
- self.assertEqual(3.0, v2.values().eval())
+ self.assertAllClose(1.0, self.evaluate(one))
+ self.assertAllClose([2.0, 2.0, 2.0], self.evaluate(twos))
+ self.assertEqual(b"k1", self.evaluate(v2.keys()))
+ self.assertEqual(3.0, self.evaluate(v2.values()))
def testVarListShouldBeEmptyInDeferredBuild(self):
with ops_lib.Graph().as_default():
@@ -664,10 +664,10 @@
# Saver with no arg, defaults to 'all variables'.
save = saver_module.Saver()
save.restore(sess, save_path)
- self.assertAllClose(1.0, one.eval())
- self.assertAllClose([2.0, 2.0, 2.0], twos.eval())
+ self.assertAllClose(1.0, self.evaluate(one))
+ self.assertAllClose([2.0, 2.0, 2.0], self.evaluate(twos))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testReshape(self):
save_path = os.path.join(self.get_temp_dir(), "variables_reshape")
with session.Session("", graph=ops_lib.Graph()) as sess:
@@ -691,7 +691,8 @@
var = variables.VariableV1([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
save = saver_module.Saver(reshape=True)
save.restore(sess, save_path)
- self.assertAllClose([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], var.eval())
+ self.assertAllClose([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
+ self.evaluate(var))
@test_util.run_in_graph_and_eager_modes
def testSaveWithGlobalStep(self, pad_step_number=False):
@@ -726,7 +727,6 @@
def testSaveWithGlobalStepWithPadding(self):
self.testSaveWithGlobalStep(pad_step_number=True)
- @test_util.run_deprecated_v1
def testSaveToNonexistingPath(self):
file_io.write_string_to_file(
os.path.join(self.get_temp_dir(), "actually_a_file"), "")
@@ -753,8 +753,8 @@
self.evaluate(init_all_op)
# Check that the parameter nodes have been initialized.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
# Save the graph.
save.save(sess, save_path)
@@ -763,13 +763,12 @@
# Restore the saved values in the parameter nodes.
save.restore(sess, save_path)
# Check that the parameter nodes have been restored.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
except ValueError as exc:
error_msg_template = "Parent directory of {} doesn't exist, can't save."
self.assertEqual(error_msg_template.format(save_path), str(exc))
- @test_util.run_deprecated_v1
def testSaveToURI(self):
# ParseURI functions don't work on Windows yet.
# TODO(jhseu): Remove this check when it works.
@@ -789,8 +788,8 @@
self.evaluate(init_all_op)
# Check that the parameter nodes have been initialized.
- self.assertEqual(10.0, v0.eval())
- self.assertEqual(20.0, v1.eval())
+ self.assertEqual(10.0, self.evaluate(v0))
+ self.assertEqual(20.0, self.evaluate(v1))
save.save(sess, save_path)
def testSaveRestoreAndValidateVariableDtype(self):
@@ -835,7 +834,7 @@
orig_vars = _model()
self.evaluate(variables.global_variables_initializer())
save = saver_module.Saver(max_to_keep=1)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
save.save(sess, save_dir)
orig_vals = self.evaluate(orig_vars)
@@ -882,7 +881,7 @@
},
write_version=self._WRITE_VERSION,
sharded=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
t0.insert("k1", 30.0).run()
t1.insert("k2", 40.0).run()
val = save.save(sess, save_path)
@@ -908,15 +907,15 @@
},
write_version=self._WRITE_VERSION,
sharded=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
t0.insert("k11", 33.0).run()
- self.assertEqual(111, v0.eval())
- self.assertEqual(b"k11", t0.keys().eval())
- self.assertEqual(33.0, t0.values().eval())
+ self.assertEqual(111, self.evaluate(v0))
+ self.assertEqual(b"k11", self.evaluate(t0.keys()))
+ self.assertEqual(33.0, self.evaluate(t0.values()))
save.restore(sess, save_path + "-00000-of-00002")
- self.assertEqual(10, v0.eval())
- self.assertEqual(b"k1", t0.keys().eval())
- self.assertEqual(30.0, t0.values().eval())
+ self.assertEqual(10, self.evaluate(v0))
+ self.assertEqual(b"k1", self.evaluate(t0.keys()))
+ self.assertEqual(30.0, self.evaluate(t0.values()))
# Restore different ops from shard 1 of the saved files.
with session.Session(
@@ -932,15 +931,15 @@
},
write_version=self._WRITE_VERSION,
sharded=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
t1.insert("k22", 44.0).run()
- self.assertEqual(222, v1.eval())
- self.assertEqual(b"k22", t1.keys().eval())
- self.assertEqual(44.0, t1.values().eval())
+ self.assertEqual(222, self.evaluate(v1))
+ self.assertEqual(b"k22", self.evaluate(t1.keys()))
+ self.assertEqual(44.0, self.evaluate(t1.values()))
save.restore(sess, save_path + "-00001-of-00002")
- self.assertEqual(20, v1.eval())
- self.assertEqual(b"k2", t1.keys().eval())
- self.assertEqual(40.0, t1.values().eval())
+ self.assertEqual(20, self.evaluate(v1))
+ self.assertEqual(b"k2", self.evaluate(t1.keys()))
+ self.assertEqual(40.0, self.evaluate(t1.values()))
# Now try a restore with the sharded filename.
with session.Session(
@@ -961,26 +960,26 @@
},
write_version=self._WRITE_VERSION,
sharded=True)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
t0.insert("k11", 33.0).run()
t1.insert("k22", 44.0).run()
- self.assertEqual(111, v0.eval())
- self.assertEqual(222, v1.eval())
- self.assertEqual(b"k11", t0.keys().eval())
- self.assertEqual(33.0, t0.values().eval())
- self.assertEqual(b"k22", t1.keys().eval())
- self.assertEqual(44.0, t1.values().eval())
+ self.assertEqual(111, self.evaluate(v0))
+ self.assertEqual(222, self.evaluate(v1))
+ self.assertEqual(b"k11", self.evaluate(t0.keys()))
+ self.assertEqual(33.0, self.evaluate(t0.values()))
+ self.assertEqual(b"k22", self.evaluate(t1.keys()))
+ self.assertEqual(44.0, self.evaluate(t1.values()))
save_path = os.path.join(self.get_temp_dir(), "sharded_basics")
if save._write_version is saver_pb2.SaverDef.V1:
save.restore(sess, save_path + "-?????-of-?????")
else:
save.restore(sess, save_path)
- self.assertEqual(10, v0.eval())
- self.assertEqual(20, v1.eval())
- self.assertEqual(b"k1", t0.keys().eval())
- self.assertEqual(30.0, t0.values().eval())
- self.assertEqual(b"k2", t1.keys().eval())
- self.assertEqual(40.0, t1.values().eval())
+ self.assertEqual(10, self.evaluate(v0))
+ self.assertEqual(20, self.evaluate(v1))
+ self.assertEqual(b"k1", self.evaluate(t0.keys()))
+ self.assertEqual(30.0, self.evaluate(t0.values()))
+ self.assertEqual(b"k2", self.evaluate(t1.keys()))
+ self.assertEqual(40.0, self.evaluate(t1.values()))
if save._write_version is saver_pb2.SaverDef.V1:
self.assertEqual(
@@ -1028,7 +1027,7 @@
else:
vs = [variables.VariableV1(rnd, name=var_name)]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
if call_saver_with_dict:
saver = saver_module.Saver({var_name: vs[0]})
else:
@@ -1056,7 +1055,7 @@
name=var_name)
]
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
if call_saver_with_dict:
saver = saver_module.Saver({
var_name: new_vs[0]
@@ -1203,7 +1202,7 @@
with self.cached_session() as sess:
v = variables.VariableV1(10.0, name="v")
save = saver_module.Saver({"v": v}, max_to_keep=2)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual([], save.last_checkpoints)
s1 = save.save(sess, os.path.join(save_dir, "s1"))
@@ -1388,7 +1387,7 @@
"v0": v0,
"v1": v1
}, sharded=True, max_to_keep=2)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual([], save.last_checkpoints)
s1 = save.save(sess, os.path.join(save_dir, "s1"))
@@ -1434,14 +1433,13 @@
self.assertTrue(
gfile.Exists(checkpoint_management.meta_graph_filename(s3)))
- @test_util.run_deprecated_v1
def testNoMaxToKeep(self):
save_dir = self._get_test_dir("no_max_to_keep")
save_dir2 = self._get_test_dir("max_to_keep_0")
with self.cached_session() as sess:
v = variables.VariableV1(10.0, name="v")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Test max_to_keep being None.
save = saver_module.Saver({"v": v}, max_to_keep=None)
@@ -1463,14 +1461,13 @@
self.assertEqual([], save2.last_checkpoints)
self.assertTrue(checkpoint_management.checkpoint_exists(s2))
- @test_util.run_deprecated_v1
def testNoMetaGraph(self):
save_dir = self._get_test_dir("no_meta_graph")
with self.cached_session() as sess:
v = variables.VariableV1(10.0, name="v")
save = saver_module.Saver({"v": v})
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
s1 = save.save(sess, os.path.join(save_dir, "s1"), write_meta_graph=False)
self.assertTrue(checkpoint_management.checkpoint_exists(s1))
@@ -1487,7 +1484,6 @@
@test_util.run_in_graph_and_eager_modes
@test.mock.patch.object(saver_module, "time")
- @test_util.run_deprecated_v1
def testNonSharded(self, mock_time):
save_dir = self._get_test_dir("keep_checkpoint_every_n_hours")
@@ -1607,7 +1603,6 @@
self.assertEqual(20.0, self.evaluate(v1))
@test_util.run_in_graph_and_eager_modes
- @test_util.run_deprecated_v1
def testNonReshapeResourceVariable(self):
self._testNonReshape(resource_variable_ops.ResourceVariable)
@@ -1622,7 +1617,7 @@
gfile.MakeDirs(test_dir)
return test_dir
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testAddCollectionDef(self):
test_dir = self._get_test_dir("good_collection")
filename = os.path.join(test_dir, "metafile")
@@ -1714,7 +1709,7 @@
saver1 = saver_module.Saver({"v1": v1}, name="saver1")
ops_lib.add_to_collection("savers", saver0)
ops_lib.add_to_collection("savers", saver1)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Saves to different checkpoints.
saver0.save(sess, saver0_ckpt)
saver1.save(sess, saver1_ckpt)
@@ -1760,7 +1755,8 @@
new_saver0.restore(sess, saver0_ckpt)
v0 = sess.graph.get_tensor_by_name("v0:0")
v1 = sess.graph.get_tensor_by_name("v1:0")
- self.assertAllEqual([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], v0.eval())
+ self.assertAllEqual([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
+ self.evaluate(v0))
self.assertEqual([3, 2], v0.get_shape())
self.assertEqual([], v1.get_shape())
with self.assertRaisesWithPredicateMatch(
@@ -1770,15 +1766,15 @@
new_saver1 = savers[1]
new_saver1.restore(sess, saver1_ckpt)
v1 = sess.graph.get_tensor_by_name("v1:0")
- self.assertEqual(11.0, v1.eval())
+ self.assertEqual(11.0, self.evaluate(v1))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testMultiSaverCollection(self):
test_dir = self._get_test_dir("saver_collection")
self._testMultiSaverCollectionSave(test_dir)
self._testMultiSaverCollectionRestore(test_dir)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testClearExtraneousSavers(self):
test_dir = self._get_test_dir("clear_extraneous_savers")
filename = os.path.join(test_dir, "metafile")
@@ -1794,7 +1790,7 @@
saver1 = saver_module.Saver({"v1": v1}, name="saver1")
ops_lib.add_to_collection("savers", saver0)
ops_lib.add_to_collection("savers", saver1)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Saves to different checkpoints.
saver0.save(sess, saver0_ckpt)
@@ -1866,7 +1862,7 @@
lambda e: "does not exist"):
saver_module.import_meta_graph(filename)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSliceVariable(self):
test_dir = self._get_test_dir("slice_saver")
filename = os.path.join(test_dir, "metafile")
@@ -1878,7 +1874,7 @@
# The names are different and will work.
slice_saver = saver_module.Saver({"first": v1, "second": v2})
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
# Exports to meta_graph
meta_graph_def = slice_saver.export_meta_graph(filename)
@@ -2093,7 +2089,6 @@
return i + 1, x + r
self._testWhileLoopAndGradientSerDes(body)
- @test_util.run_deprecated_v1
def testNestedControlFlowSerDes(self):
# Test while loop in a cond in a while loop.
# pylint: disable=g-long-lambda
@@ -2122,7 +2117,7 @@
lambda: math_ops.multiply(x, -1.0))))
# pylint: enable=g-long-lambda
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testStrippedOpListDef(self):
with self.cached_session():
# Creates a graph.
@@ -2745,7 +2740,7 @@
graph.add_to_collection(ops_lib.GraphKeys.SAVERS, saver2)
with self.session(graph=graph) as sess:
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
saver1.save(sess, saver1_ckpt, write_state=False)
saver2.save(sess, saver2_ckpt, write_state=False)
@@ -2762,7 +2757,7 @@
with self.session(graph=graph1) as sess:
saver_list1[0].restore(sess, saver1_ckpt)
- self.assertEqual(1.0, var_dict1["variable1:0"].eval())
+ self.assertEqual(1.0, self.evaluate(var_dict1["variable1:0"]))
graph2 = ops_lib.Graph()
var_dict2 = meta_graph.copy_scoped_meta_graph(
@@ -2777,7 +2772,7 @@
with self.session(graph=graph2) as sess:
saver_list2[0].restore(sess, saver2_ckpt)
- self.assertEqual(2.0, var_dict2["variable2:0"].eval())
+ self.assertEqual(2.0, self.evaluate(var_dict2["variable2:0"]))
class _OwnsAVariableSimple(checkpointable_base.CheckpointableBase):
@@ -2988,7 +2983,7 @@
# exception" block in Python 3.
self.assertNotIn("NewCheckpointReader", cs.exception.message)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testGraphChangedForRestoreErrorRaised(self):
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
@@ -3010,7 +3005,6 @@
"a mismatch between the current graph and the graph"):
a_saver.restore(sess=sess, save_path=save_path)
- @test_util.run_deprecated_v1
def testLoadFromObjectBasedGraph(self):
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
diff --git a/tensorflow/python/training/saving/BUILD b/tensorflow/python/training/saving/BUILD
new file mode 100644
index 0000000..67ccd59
--- /dev/null
+++ b/tensorflow/python/training/saving/BUILD
@@ -0,0 +1,55 @@
+# Description:
+# Low-level utilities for reading and writing checkpoints.
+
+package(
+ default_visibility = [
+ "//tensorflow:internal",
+ ],
+)
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(["LICENSE"])
+
+load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+
+py_library(
+ name = "functional_saver",
+ srcs = ["functional_saver.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":saveable_object",
+ ":saveable_object_util",
+ "//tensorflow/python/eager:def_function",
+ ],
+)
+
+cuda_py_test(
+ name = "functional_saver_test",
+ size = "medium",
+ srcs = [
+ "functional_saver_test.py",
+ ],
+ additional_deps = [
+ ":functional_saver",
+ "//tensorflow/python/eager:test",
+ ],
+)
+
+py_library(
+ name = "saveable_object",
+ srcs = ["saveable_object.py"],
+ srcs_version = "PY2AND3",
+)
+
+py_library(
+ name = "saveable_object_util",
+ srcs = ["saveable_object_util.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow/python:resource_variable_ops",
+ "//tensorflow/python:variables",
+ "//tensorflow/python/training/checkpointable:base",
+ "@six_archive//:six",
+ ],
+)
diff --git a/tensorflow/python/training/saving/functional_saver.py b/tensorflow/python/training/saving/functional_saver.py
new file mode 100644
index 0000000..7eed333
--- /dev/null
+++ b/tensorflow/python/training/saving/functional_saver.py
@@ -0,0 +1,101 @@
+# 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.
+# ==============================================================================
+"""Saves and restore variables inside traced @tf.functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_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 io_ops
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.training.saving import saveable_object_util
+
+
+class Saver(object):
+ """A minimal utility class for saving and restoring checkpoints.
+
+ Note that this is a low-level utility which stores Tensors in the keys
+ specified by `SaveableObject`s. Higher-level utilities for object-based
+ checkpointing are built on top of it.
+ """
+
+ def __init__(self, saveable_objects):
+ """Specify a list of `SaveableObject`s to save and restore.
+
+ Args:
+ saveable_objects: A list of `SaveableObject`s.
+ """
+ saveable_objects = list(saveable_objects)
+ for saveable in saveable_objects:
+ if not isinstance(saveable, saveable_object.SaveableObject):
+ raise ValueError(
+ "Saver expected a list of SaveableObjects, got %s." % (saveable,))
+ self._saveable_objects = saveable_objects
+
+ # TODO(b/120569892): Use tf.function here
+ def save(self, file_prefix):
+ """Save the saveable objects to a checkpoint with `file_prefix`.
+
+ Args:
+ file_prefix: A string or scalar string Tensor containing the prefix to
+ save under.
+ Returns:
+ A scalar string Tensor containing `file_prefix` with control dependencies
+ on the save ops.
+ """
+ tensor_names = []
+ tensors = []
+ tensor_slices = []
+ for saveable in self._saveable_objects:
+ for spec in saveable.specs:
+ tensor_names.append(spec.name)
+ tensors.append(spec.tensor)
+ tensor_slices.append(spec.slice_spec)
+ with ops.control_dependencies(
+ [io_ops.save_v2(file_prefix, tensor_names, tensor_slices, tensors)]):
+ return array_ops.identity(file_prefix)
+
+ # TODO(b/120569892): Use tf.function here
+ def restore(self, file_prefix):
+ """Restore the saveable objects from a checkpoint with `file_prefix`.
+
+ Args:
+ file_prefix: A string or scalar string Tensor containing the prefix for
+ files to read from.
+
+ Returns:
+ An operation which restores the `Saver`'s `SaveableObject`s when run, or
+ None if executing eagerly.
+ """
+ restore_ops = []
+ for saveable in self._saveable_objects:
+ if saveable.device:
+ device = saveable_object_util.set_cpu0(saveable.device)
+ else:
+ device = None
+ with ops.device(device):
+ tensors = []
+ for spec in saveable.specs:
+ tensors.append(
+ io_ops.restore_v2(
+ file_prefix,
+ [spec.name],
+ [spec.slice_spec],
+ [spec.dtype])[0])
+ restore_ops.append(saveable.restore(tensors, restored_shapes=None))
+ return control_flow_ops.group(restore_ops)
diff --git a/tensorflow/python/training/saving/functional_saver_test.py b/tensorflow/python/training/saving/functional_saver_test.py
new file mode 100644
index 0000000..4000225
--- /dev/null
+++ b/tensorflow/python/training/saving/functional_saver_test.py
@@ -0,0 +1,50 @@
+# 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 the functional saver."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.training.saving import functional_saver
+from tensorflow.python.training.saving import saveable_object_util
+
+
+class SaverTest(test.TestCase):
+
+ def test_resource_variable(self):
+ v1 = resource_variable_ops.ResourceVariable(2.)
+ saver = functional_saver.Saver(
+ saveable_object_util.saveable_objects_for_op(v1, "x"))
+ prefix = os.path.join(self.get_temp_dir(), "ckpt")
+ save_path = saver.save(constant_op.constant(prefix))
+ v1.assign(1.)
+ saver.restore(save_path)
+ self.assertEqual(2., self.evaluate(v1))
+
+ v2 = resource_variable_ops.ResourceVariable(3.)
+ second_saver = functional_saver.Saver(
+ saveable_object_util.saveable_objects_for_op(v2, "x"))
+ second_saver.restore(save_path)
+ self.assertEqual(2., self.evaluate(v2))
+
+
+if __name__ == "__main__":
+ test.main()
diff --git a/tensorflow/python/training/saveable_object.py b/tensorflow/python/training/saving/saveable_object.py
similarity index 100%
rename from tensorflow/python/training/saveable_object.py
rename to tensorflow/python/training/saving/saveable_object.py
diff --git a/tensorflow/python/training/saving/saveable_object_util.py b/tensorflow/python/training/saving/saveable_object_util.py
new file mode 100644
index 0000000..fa88d2c
--- /dev/null
+++ b/tensorflow/python/training/saving/saveable_object_util.py
@@ -0,0 +1,340 @@
+# 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.
+# ==============================================================================
+"""Utilities for working with and creating SaveableObjects."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import six
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.training.checkpointable import base as checkpointable
+from tensorflow.python.training.saving import saveable_object
+
+
+# Op names which identify variable reads which should be saved.
+_VARIABLE_OPS = set(["Variable",
+ "VariableV2",
+ "AutoReloadVariable",
+ "VarHandleOp",
+ "ReadVariableOp"])
+
+
+def set_cpu0(device_string):
+ """Creates a new device string based on `device_string` but using /CPU:0.
+
+ If the device is already on /CPU:0, this is a no-op.
+
+ Args:
+ device_string: A device string.
+
+ Returns:
+ A device string.
+ """
+ parsed_device = pydev.DeviceSpec.from_string(device_string)
+ parsed_device.device_type = "CPU"
+ parsed_device.device_index = 0
+ return parsed_device.to_string()
+
+
+class ReferenceVariableSaveable(saveable_object.SaveableObject):
+ """SaveableObject implementation that handles reference variables."""
+
+ def __init__(self, var, slice_spec, name):
+ spec = saveable_object.SaveSpec(var, slice_spec, name, dtype=var.dtype)
+ super(ReferenceVariableSaveable, self).__init__(var, [spec], name)
+
+ def restore(self, restored_tensors, restored_shapes):
+ restored_tensor = restored_tensors[0]
+ if restored_shapes is not None:
+ restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
+ return state_ops.assign(
+ self.op,
+ restored_tensor,
+ validate_shape=restored_shapes is None and
+ self.op.get_shape().is_fully_defined())
+
+
+class ResourceVariableSaveable(saveable_object.SaveableObject):
+ """SaveableObject implementation that handles ResourceVariables."""
+
+ def __init__(self, var, slice_spec, name):
+ self._var_device = var.device
+ self._var_shape = var.shape
+ if isinstance(var, ops.Tensor):
+ self.handle_op = var.op.inputs[0]
+ tensor = var
+ elif isinstance(var, resource_variable_ops.ResourceVariable):
+
+ def _read_variable_closure(v):
+ def f():
+ with ops.device(v.device):
+ x = v.read_value()
+ # To allow variables placed on non-CPU devices to be checkpointed,
+ # we copy them to CPU on the same machine first.
+ with ops.device("/device:CPU:0"):
+ return array_ops.identity(x)
+ return f
+
+ self.handle_op = var.handle
+ tensor = _read_variable_closure(var)
+ else:
+ raise ValueError(
+ "Saveable is neither a resource variable nor a read operation."
+ " Got: %s" % repr(var))
+ spec = saveable_object.SaveSpec(tensor, slice_spec, name,
+ dtype=var.dtype)
+ super(ResourceVariableSaveable, self).__init__(var, [spec], name)
+
+ def restore(self, restored_tensors, restored_shapes):
+ restored_tensor = restored_tensors[0]
+ if restored_shapes is not None:
+ restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
+ # Copy the restored tensor to the variable's device.
+ with ops.device(self._var_device):
+ restored_tensor = array_ops.identity(restored_tensor)
+ return resource_variable_ops.shape_safe_assign_variable_handle(
+ self.handle_op, self._var_shape, restored_tensor)
+
+
+def _tensor_comes_from_variable(v):
+ return isinstance(v, ops.Tensor) and v.op.type in _VARIABLE_OPS
+
+
+def saveable_objects_for_op(op, name):
+ """Create `SaveableObject`s from an operation.
+
+ Args:
+ op: A variable, operation, or SaveableObject to coerce into a
+ SaveableObject.
+ name: A string name for the SaveableObject.
+
+ Yields:
+ `SaveableObject`s which together save/restore `op`.
+
+ Raises:
+ TypeError: If `name` is not a string.
+ ValueError: For operations with no known conversion to SaveableObject.
+ """
+ if not isinstance(name, six.string_types):
+ raise TypeError(
+ "names_to_saveables must be a dict mapping string names to "
+ "checkpointable operations. Name is not a string: %s" % name)
+ if isinstance(op, saveable_object.SaveableObject):
+ yield op
+ elif isinstance(op, (list, tuple, variables.PartitionedVariable)):
+ if isinstance(op, variables.PartitionedVariable):
+ op = list(op)
+ # A set of slices.
+ slice_name = None
+ # pylint: disable=protected-access
+ for variable in op:
+ if not isinstance(variable, variables.Variable):
+ raise ValueError("Slices must all be Variables: %s" % variable)
+ if not variable._save_slice_info:
+ raise ValueError("Slices must all be slices: %s" % variable)
+ if slice_name is None:
+ slice_name = variable._save_slice_info.full_name
+ elif slice_name != variable._save_slice_info.full_name:
+ raise ValueError(
+ "Slices must all be from the same tensor: %s != %s" %
+ (slice_name, variable._save_slice_info.full_name))
+ if variable.op.type in ["Variable", "VariableV2",
+ "AutoReloadVariable"]:
+ yield ReferenceVariableSaveable(
+ variable, variable._save_slice_info.spec, name)
+ else:
+ yield ResourceVariableSaveable(
+ variable, variable._save_slice_info.spec, name)
+ # pylint: enable=protected-access
+ elif isinstance(op, checkpointable.CheckpointableBase) and not isinstance(
+ op, variables.Variable):
+ # pylint: disable=protected-access
+ for attr, factory in op._gather_saveables_for_checkpoint().items():
+ if attr == checkpointable.VARIABLE_VALUE_KEY:
+ # Keep original name for classes masquerading as variables.
+ full_name = name
+ else:
+ full_name = name + "_" + attr
+ op = (factory(full_name) if callable(factory) else factory)
+ for op in saveable_objects_for_op(op, op.name):
+ yield op
+ # pylint: enable=protected-access
+ else:
+ # A variable or tensor.
+ if isinstance(op, resource_variable_ops.ResourceVariable):
+ # pylint: disable=protected-access
+ if op._in_graph_mode:
+ variable = op._graph_element
+ else:
+ variable = op
+ # pylint: enable=protected-access
+ yield ResourceVariableSaveable(variable, "", name)
+ else:
+ with ops.init_scope():
+ if context.executing_eagerly():
+ raise ValueError("Can only save/restore ResourceVariables when "
+ "executing eagerly, got type: %s." % type(op))
+
+ variable = ops.internal_convert_to_tensor(op, as_ref=True)
+ if not _tensor_comes_from_variable(variable):
+ raise TypeError("names_to_saveables must be a dict mapping string "
+ "names to Tensors/Variables. Not a variable: %s" %
+ variable)
+ if variable.op.type in ["Variable", "VariableV2",
+ "AutoReloadVariable"]:
+ yield ReferenceVariableSaveable(variable, "", name)
+ else:
+ yield ResourceVariableSaveable(
+ variable, "", name)
+
+
+def op_list_to_dict(op_list, convert_variable_to_tensor=True):
+ """Create a dictionary of names to operation lists.
+
+ Args:
+ op_list: A list, tuple, or set of Variables or SaveableObjects.
+ convert_variable_to_tensor: Whether or not to convert single Variables
+ with no slice info into Tensors.
+
+ Returns:
+ A dictionary of names to the operations that must be saved under
+ that name. Variables with save_slice_info are grouped together under the
+ same key in no particular order.
+
+ Raises:
+ TypeError: If the type of op_list or its elements is not supported.
+ ValueError: If at least two saveables share the same name.
+ """
+ if not isinstance(op_list, (list, tuple, set)):
+ raise TypeError("Variables to save should be passed in a dict or a "
+ "list: %s" % op_list)
+ # When ResourceVariables are converted to Tensors, read ops are added to the
+ # graph. Sorting the op_list ensures that the resulting graph is always
+ # constructed in a deterministic way:
+ op_list = sorted(op_list, key=lambda x: x.name)
+ names_to_saveables = {}
+ # pylint: disable=protected-access
+ for var in op_list:
+ if isinstance(var, saveable_object.SaveableObject):
+ names_to_saveables[var.name] = var
+ elif isinstance(var, variables.PartitionedVariable):
+ if var.name in names_to_saveables:
+ raise ValueError("At least two variables have the same name: %s" %
+ var.name)
+ names_to_saveables[var.name] = var
+ elif isinstance(var, variables.Variable) and var._save_slice_info:
+ name = var._save_slice_info.full_name
+ if name in names_to_saveables:
+ if not isinstance(names_to_saveables[name], list):
+ raise ValueError("Mixing slices and non-slices with the same name: "
+ "%s" % name)
+ names_to_saveables[name].append(var)
+ else:
+ names_to_saveables[name] = [var]
+ elif (isinstance(var, checkpointable.CheckpointableBase)
+ and not isinstance(var, variables.Variable)):
+ checkpointable_saveables = [
+ (factory() if callable(factory) else factory)
+ for factory in var._gather_saveables_for_checkpoint().values()]
+ names_to_saveables.update(
+ op_list_to_dict(checkpointable_saveables))
+ else:
+ if context.executing_eagerly():
+ if not isinstance(var, resource_variable_ops.ResourceVariable):
+ raise ValueError(
+ "Can only save/restore ResourceVariables when eager execution "
+ "is enabled, type: %s." % type(var))
+ set_var = names_to_saveables.setdefault(var._shared_name, var)
+ if set_var is not var:
+ raise ValueError(
+ ("Two different ResourceVariable objects with the same "
+ "shared_name '%s' were passed to the Saver. This likely means "
+ "that they were created in different Graphs or isolation "
+ "contexts, and may not be checkpointed together.") %
+ (var._shared_name,))
+ else:
+ if convert_variable_to_tensor:
+ if isinstance(var, resource_variable_ops.ResourceVariable):
+ var = var._graph_element # pylint: disable=protected-access
+ else:
+ var = ops.internal_convert_to_tensor(var, as_ref=True)
+ if not _tensor_comes_from_variable(var):
+ raise TypeError("Variable to save is not a Variable: %s" % var)
+ if var.op.type == "ReadVariableOp":
+ name = var.op.inputs[0].op.name
+ else:
+ name = var.op.name
+ if name in names_to_saveables:
+ raise ValueError("At least two variables have the same name: %s" %
+ name)
+ names_to_saveables[name] = var
+
+ # pylint: enable=protected-access
+ return names_to_saveables
+
+
+def _add_saveable(saveables, seen_ops, saveable):
+ """Adds the saveable to the saveables list.
+
+ Args:
+ saveables: List to append the SaveableObject to.
+ seen_ops: Set of the ops of the saveables already processed. Used to
+ check that each saveable is only saved once.
+ saveable: The saveable.
+
+ Raises:
+ ValueError: If the saveable has already been processed.
+ """
+ if saveable.op in seen_ops:
+ raise ValueError("The same saveable will be restored with two names: %s" %
+ saveable.name)
+ saveables.append(saveable)
+ seen_ops.add(saveable.op)
+
+
+def validate_and_slice_inputs(names_to_saveables):
+ """Returns the variables and names that will be used for a Saver.
+
+ Args:
+ names_to_saveables: A dict (k, v) where k is the name of an operation and
+ v is an operation to save or a BaseSaverBuilder.Saver.
+
+ Returns:
+ A list of SaveableObjects.
+
+ Raises:
+ TypeError: If any of the keys are not strings or any of the
+ values are not one of Tensor or Variable or a checkpointable operation.
+ ValueError: If the same operation is given in more than one value
+ (this also applies to slices of SlicedVariables).
+ """
+ if not isinstance(names_to_saveables, dict):
+ names_to_saveables = op_list_to_dict(names_to_saveables)
+
+ saveables = []
+ seen_ops = set()
+ for name, op in sorted(names_to_saveables.items(),
+ # Avoid comparing ops, sort only by name.
+ key=lambda x: x[0]):
+ for converted_saveable_object in saveable_objects_for_op(op, name):
+ _add_saveable(saveables, seen_ops, converted_saveable_object)
+ return saveables
diff --git a/tensorflow/python/training/server_lib.py b/tensorflow/python/training/server_lib.py
index 302ca2d..b3c21d5 100644
--- a/tensorflow/python/training/server_lib.py
+++ b/tensorflow/python/training/server_lib.py
@@ -23,6 +23,7 @@
from tensorflow.python import pywrap_tensorflow as c_api
from tensorflow.python.framework import errors
from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@@ -93,7 +94,8 @@
return server_def
-@tf_export("train.Server")
+@tf_export("distribute.Server", v1=["distribute.Server", "train.Server"])
+@deprecation.deprecated_endpoints("train.Server")
class Server(object):
"""An in-process TensorFlow server, for use in distributed training.
diff --git a/tensorflow/python/training/server_lib_same_variables_no_clear_test.py b/tensorflow/python/training/server_lib_same_variables_no_clear_test.py
index 1b2d588..ff3fab9 100644
--- a/tensorflow/python/training/server_lib_same_variables_no_clear_test.py
+++ b/tensorflow/python/training/server_lib_same_variables_no_clear_test.py
@@ -34,7 +34,7 @@
# TODO(b/34465411): Starting multiple servers with different configurations
# in the same test is flaky. Move this test case back into
# "server_lib_test.py" when this is no longer the case.
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSameVariablesNoClear(self):
server = server_lib.Server.create_local_server()
diff --git a/tensorflow/python/training/server_lib_test.py b/tensorflow/python/training/server_lib_test.py
index 323e94c..92cdc1c 100644
--- a/tensorflow/python/training/server_lib_test.py
+++ b/tensorflow/python/training/server_lib_test.py
@@ -29,6 +29,7 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import math_ops
@@ -55,6 +56,7 @@
self.assertAllEqual([[4]], sess.run(e))
# TODO(mrry): Add `server.stop()` and `server.join()` when these work.
+ @test_util.run_v1_only("b/120545219")
def testMultipleSessions(self):
server = self._cached_server
@@ -73,6 +75,7 @@
# TODO(mrry): Add `server.stop()` and `server.join()` when these work.
# Verifies various reset failures.
+ @test_util.run_v1_only("b/120545219")
def testResetFails(self):
# Creates variable with container name.
with ops.container("test0"):
@@ -146,6 +149,7 @@
self.assertEqual(0.5, min_val)
self.assertEqual(0.5, max_val)
+ @test_util.run_v1_only("b/120545219")
def testCloseCancelsBlockingOperation(self):
server = self._cached_server
sess = session.Session(server.target, config=self._useRPCConfig())
@@ -207,6 +211,7 @@
"local": ["localhost"]
}, job_name="local", task_index=0)
+ @test_util.run_v1_only("b/120545219")
def testTimeoutRaisesException(self):
server = self._cached_server
q = data_flow_ops.FIFOQueue(1, [dtypes.float32])
@@ -241,6 +246,7 @@
queue_runner_impl.start_queue_runners(sess)
sess.run(var.assign(3.0))
+ @test_util.run_v1_only("b/120545219")
def testIsolateSessionState(self):
server = self._cached_server
@@ -296,6 +302,7 @@
self.assertAllEqual(37, isolate_sess_0.run(v))
self.assertAllEqual([19, 86], isolate_sess_1.run(v))
+ @test_util.run_v1_only("b/120545219")
def testShapeChangingIsolateState(self):
server = self._cached_server
sharing_config = config_pb2.ConfigProto(isolate_session_state=False)
diff --git a/tensorflow/python/training/session_manager.py b/tensorflow/python/training/session_manager.py
index 1465863..0f68fcf 100644
--- a/tensorflow/python/training/session_manager.py
+++ b/tensorflow/python/training/session_manager.py
@@ -25,7 +25,6 @@
from tensorflow.python.framework import ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import checkpoint_management
-from tensorflow.python.training import distribution_strategy_context
from tensorflow.python.util.tf_export import tf_export
@@ -183,12 +182,6 @@
"""
self._target = master
sess = session.Session(self._target, graph=self._graph, config=config)
- # TODO(jhseu): Delete once tpu.initialize_system() goes away.
- initialize_ops = (
- distribution_strategy_context.get_distribution_strategy().initialize()
- )
- if initialize_ops:
- sess.run(initialize_ops)
if checkpoint_dir and checkpoint_filename_with_path:
raise ValueError("Can not provide both checkpoint_dir and "
diff --git a/tensorflow/python/training/session_manager_test.py b/tensorflow/python/training/session_manager_test.py
index 4294ffa..c9a0c56 100644
--- a/tensorflow/python/training/session_manager_test.py
+++ b/tensorflow/python/training/session_manager_test.py
@@ -69,7 +69,7 @@
"", init_fn=lambda sess: sess.run(v.initializer))
self.assertAllClose([125], sess.run(v))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPrepareSessionFails(self):
checkpoint_dir = os.path.join(self.get_temp_dir(), "prepare_session")
checkpoint_dir2 = os.path.join(self.get_temp_dir(), "prepare_session2")
@@ -154,7 +154,7 @@
sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
self.assertEquals(1, sess.run(v))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRecoverSession(self):
# Create a checkpoint.
checkpoint_dir = os.path.join(self.get_temp_dir(), "recover_session")
@@ -187,6 +187,7 @@
checkpoint_filename_with_path=checkpoint_management.latest_checkpoint(
checkpoint_dir))
+ @test_util.run_v1_only("b/120545219")
def testWaitForSessionReturnsNoneAfterTimeout(self):
with ops.Graph().as_default():
variables.VariableV1(1, name="v")
@@ -209,7 +210,7 @@
variables.global_variables()),
local_init_op=None)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRecoverSessionWithReadyForLocalInitOp(self):
# Create a checkpoint.
checkpoint_dir = os.path.join(self.get_temp_dir(),
@@ -263,7 +264,7 @@
self.assertEquals(1, sess.run(v))
self.assertEquals(1, sess.run(w))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRecoverSessionWithReadyForLocalInitOpFailsToReadyLocal(self):
# We use ready_for_local_init_op=tf.report_uninitialized_variables(),
# which causes recover_session to not run local_init_op, and to return
@@ -320,7 +321,7 @@
sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
self.assertEquals(1, sess.run(v))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRecoverSessionNoChkptStillRunsLocalInitOp(self):
# This test checks for backwards compatibility.
# In particular, we continue to ensure that recover_session will execute
@@ -349,7 +350,7 @@
sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
self.assertEquals(1, sess.run(w))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRecoverSessionFailsStillRunsLocalInitOp(self):
# Create a checkpoint.
checkpoint_dir = os.path.join(
@@ -393,7 +394,7 @@
sess.graph.get_tensor_by_name("w:0")).eval(session=sess))
self.assertEquals(1, sess.run(w))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWaitForSessionLocalInit(self):
server = server_lib.Server.create_local_server()
with ops.Graph().as_default() as graph:
@@ -445,7 +446,7 @@
# because of overly restrictive ready_for_local_init_op
sm.wait_for_session("", max_wait_secs=3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testWaitForSessionInsufficientReadyForLocalInitCheck(self):
with ops.Graph().as_default() as graph:
v = variables.VariableV1(1, name="v")
@@ -463,7 +464,7 @@
"Session was not ready after waiting.*"):
sm.wait_for_session("", max_wait_secs=3)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPrepareSessionWithReadyForLocalInitOp(self):
with ops.Graph().as_default():
v = variables.VariableV1(1, name="v")
@@ -503,7 +504,7 @@
self.assertEquals(1, sess.run(w))
self.assertEquals(3, sess.run(x))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPrepareSessionWithPartialInitOp(self):
with ops.Graph().as_default():
v = variables.VariableV1(1, name="v")
@@ -570,7 +571,7 @@
self.assertEquals(1, sess.run(w_res))
self.assertEquals(3, sess.run(x_res))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPrepareSessionWithCyclicInitializer(self):
# Regression test. Previously Variable._build_initializer_expr would enter
# into an infinite recursion when the variable's initial_value involved
@@ -644,7 +645,7 @@
"Init operations did not make model ready for local_init"):
sm2.prepare_session("", init_op=None)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPrepareSessionWithInsufficientReadyForLocalInitCheck(self):
with ops.Graph().as_default():
v = variables.VariableV1(1, name="v")
@@ -697,7 +698,7 @@
"", init_fn=lambda sess: sess.run(v.initializer))
self.assertAllClose([125], sess.run(v))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testPrepareSessionFails(self):
checkpoint_dir = os.path.join(self.get_temp_dir(), "prepare_session")
checkpoint_dir2 = os.path.join(self.get_temp_dir(), "prepare_session2")
@@ -759,7 +760,7 @@
variables.is_variable_initialized(
sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRecoverSession(self):
# Create a checkpoint.
checkpoint_dir = os.path.join(self.get_temp_dir(), "recover_session")
@@ -798,6 +799,7 @@
sess.graph.get_tensor_by_name("v:0")).eval(session=sess))
self.assertEquals(1, sess.run(v))
+ @test_util.run_v1_only("b/120545219")
def testWaitForSessionReturnsNoneAfterTimeout(self):
with ops.Graph().as_default():
variables.VariableV1(1, name="v")
diff --git a/tensorflow/python/training/session_run_hook.py b/tensorflow/python/training/session_run_hook.py
index e9a61de..886ca46 100644
--- a/tensorflow/python/training/session_run_hook.py
+++ b/tensorflow/python/training/session_run_hook.py
@@ -94,7 +94,7 @@
from tensorflow.python.util.tf_export import tf_export
-@tf_export("train.SessionRunHook")
+@tf_export(v1=["train.SessionRunHook"])
class SessionRunHook(object):
"""Hook to extend calls to MonitoredSession.run()."""
diff --git a/tensorflow/python/training/slot_creator.py b/tensorflow/python/training/slot_creator.py
index d76b22a..ecf5a96 100644
--- a/tensorflow/python/training/slot_creator.py
+++ b/tensorflow/python/training/slot_creator.py
@@ -39,13 +39,13 @@
from __future__ import division
from __future__ import print_function
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.eager import context
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
-from tensorflow.python.training import distribution_strategy_context
def _create_slot_var(primary, val, scope, validate_shape, shape, dtype):
diff --git a/tensorflow/python/training/slot_creator_test.py b/tensorflow/python/training/slot_creator_test.py
index 1f26aaa..ec2eec3 100644
--- a/tensorflow/python/training/slot_creator_test.py
+++ b/tensorflow/python/training/slot_creator_test.py
@@ -32,13 +32,13 @@
class SlotCreatorTest(test.TestCase):
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCreateSlotFromVariable(self):
with self.cached_session():
v = variables.Variable([1.0, 2.5], name="var")
slot = slot_creator.create_slot(v, v.initialized_value(), name="slot")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("var/slot", slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
@@ -51,7 +51,7 @@
v = constant_op.constant([1.0, 2.5], name="const")
slot = slot_creator.create_slot(v, v * 2, name="slot")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("const/slot", slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
@@ -66,14 +66,14 @@
slot = slot_creator.create_zeros_slot(
v, name="slot", dtype=dtypes.float64)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("var/slot", slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
self.assertEqual(dtypes.float64, slot.dtype.base_dtype)
self.assertAllEqual([0.0, 0.0], self.evaluate(slot))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCreateZerosSlotFromDynamicShapedVariable(self):
with self.cached_session():
dyn_shape = constant_op.constant([2], dtype=dtypes.int32)
@@ -88,7 +88,7 @@
slot = slot_creator.create_zeros_slot(
v, name="slot", dtype=dtypes.float64)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("var/slot", slot.op.name)
self.assertEqual([2], array_ops.shape(slot).eval())
@@ -102,7 +102,7 @@
with ops.control_dependencies(None):
slot = slot_creator.create_zeros_slot(v, name="slot")
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("const/slot", slot.op.name)
self.assertEqual([2], slot.get_shape().as_list())
@@ -118,14 +118,14 @@
slot = slot_creator.create_zeros_slot(
v, name="slot", dtype=dtypes.float64)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertEqual("const/slot", slot.op.name)
self.assertEqual([2], array_ops.shape(slot).eval())
self.assertEqual(dtypes.float64, slot.dtype.base_dtype)
self.assertAllEqual([0.0, 0.0], self.evaluate(slot))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testCreateSlotFromVariableRespectsScope(self):
# See discussion on #2740.
with self.cached_session():
diff --git a/tensorflow/python/training/supervisor_test.py b/tensorflow/python/training/supervisor_test.py
index f6505ac..180ddb5 100644
--- a/tensorflow/python/training/supervisor_test.py
+++ b/tensorflow/python/training/supervisor_test.py
@@ -421,7 +421,7 @@
with self.assertRaisesRegexp(RuntimeError, "requires a summary writer"):
sv.summary_computed(sess, sess.run(summ))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testLogdirButExplicitlyNoSummaryWriter(self):
logdir = self._test_dir("explicit_no_summary_writer")
with ops.Graph().as_default():
@@ -507,7 +507,7 @@
sv = supervisor.Supervisor(logdir="", session_manager=sm)
sv.prepare_or_wait_for_session("")
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitOp(self):
logdir = self._test_dir("default_init_op")
with ops.Graph().as_default():
@@ -517,7 +517,7 @@
self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
sv.stop()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitFn(self):
logdir = self._test_dir("default_init_op")
with ops.Graph().as_default():
@@ -531,7 +531,7 @@
self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
sv.stop()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testInitOpWithFeedDict(self):
logdir = self._test_dir("feed_dict_init_op")
with ops.Graph().as_default():
@@ -545,7 +545,7 @@
self.assertAllClose([1.0, 2.0, 3.0], sess.run(v))
sv.stop()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testReadyForLocalInitOp(self):
server = server_lib.Server.create_local_server()
logdir = self._test_dir("default_ready_for_local_init_op")
@@ -588,7 +588,7 @@
sv0.stop()
sv1.stop()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testReadyForLocalInitOpRestoreFromCheckpoint(self):
server = server_lib.Server.create_local_server()
logdir = self._test_dir("ready_for_local_init_op_restore")
@@ -720,7 +720,7 @@
"Variables not initialized: w"):
sv.prepare_or_wait_for_session(server.target)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSetupFail(self):
logdir = self._test_dir("setup_fail")
with ops.Graph().as_default():
@@ -731,7 +731,7 @@
variables.VariableV1([1.0, 2.0, 3.0], name="v")
supervisor.Supervisor(logdir=logdir, is_chief=False)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testDefaultGlobalStep(self):
logdir = self._test_dir("default_global_step")
with ops.Graph().as_default():
@@ -741,7 +741,7 @@
self.assertEquals(287, sess.run(sv.global_step))
sv.stop()
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testRestoreFromMetaGraph(self):
logdir = self._test_dir("restore_from_meta_graph")
with ops.Graph().as_default():
@@ -763,7 +763,7 @@
# This test is based on the fact that the standard services start
# right away and get to run once before sv.stop() returns.
# We still sleep a bit to make the test robust.
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testStandardServicesWithoutGlobalStep(self):
logdir = self._test_dir("standard_services_without_global_step")
# Create a checkpoint.
@@ -814,7 +814,7 @@
# Same as testStandardServicesNoGlobalStep but with a global step.
# We should get a summary about the step time.
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testStandardServicesWithGlobalStep(self):
logdir = self._test_dir("standard_services_with_global_step")
# Create a checkpoint.
diff --git a/tensorflow/python/training/sync_replicas_optimizer.py b/tensorflow/python/training/sync_replicas_optimizer.py
index 172c141..cd4590d 100644
--- a/tensorflow/python/training/sync_replicas_optimizer.py
+++ b/tensorflow/python/training/sync_replicas_optimizer.py
@@ -19,6 +19,7 @@
from __future__ import print_function
from tensorflow.core.framework import types_pb2
+from tensorflow.python.distribute import distribution_strategy_context
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
@@ -259,7 +260,9 @@
# local_anchor op will be placed on this worker task by default.
local_anchor = control_flow_ops.no_op()
# Colocating local_step variable prevents it being placed on the PS.
- with ops.colocate_with(local_anchor):
+ distribution_strategy = (
+ distribution_strategy_context.get_distribution_strategy())
+ with distribution_strategy.colocate_vars_with(local_anchor):
self._local_step = variable_scope.variable(
initial_value=0,
trainable=False,
diff --git a/tensorflow/python/training/sync_replicas_optimizer_test.py b/tensorflow/python/training/sync_replicas_optimizer_test.py
index 1ef8756..428583d 100644
--- a/tensorflow/python/training/sync_replicas_optimizer_test.py
+++ b/tensorflow/python/training/sync_replicas_optimizer_test.py
@@ -22,6 +22,7 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
from tensorflow.python.framework.test_util import create_local_cluster
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
@@ -88,6 +89,7 @@
def _run(self, train_op, sess):
sess.run(train_op)
+ @test_util.run_v1_only("b/120545219")
def test2Workers(self):
num_workers = 2
replicas_to_aggregate = 2
@@ -178,6 +180,7 @@
sessions[1].run(var_1_g_1))
# 3 workers and one of them is backup.
+ @test_util.run_v1_only("b/120545219")
def test3Workers1Backup(self):
num_workers = 3
replicas_to_aggregate = 2
@@ -266,6 +269,7 @@
"apply_gradient should be called"):
hook.begin()
+ @test_util.run_v1_only("b/120545219")
def testCanCreatedBeforeMinimizeCalled(self):
"""This behavior is required to be integrated with Estimators."""
opt = training.SyncReplicasOptimizer(
@@ -278,6 +282,7 @@
opt.minimize(v, global_step=global_step)
hook.begin()
+ @test_util.run_v1_only("b/120545219")
def testFetchVariableList(self):
opt = training.SyncReplicasOptimizer(
opt=adam.AdamOptimizer(0.01),
diff --git a/tensorflow/python/training/tensorboard_logging_test.py b/tensorflow/python/training/tensorboard_logging_test.py
index 5af6a0a..5088ab0 100644
--- a/tensorflow/python/training/tensorboard_logging_test.py
+++ b/tensorflow/python/training/tensorboard_logging_test.py
@@ -25,6 +25,7 @@
import time
from tensorflow.core.util import event_pb2
+from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.summary import summary_iterator
@@ -32,6 +33,7 @@
from tensorflow.python.training import tensorboard_logging
+@test_util.run_v1_only("b/120545219")
class EventLoggingTest(test.TestCase):
def setUp(self):
@@ -85,6 +87,7 @@
(event_pb2.LogMessage.ERROR, "format")])
self.assertEqual(2, self.logged_message_count)
+ @test_util.run_v1_only("b/120545219")
def testVerbosity(self):
tensorboard_logging.set_summary_writer(self._sw)
tensorboard_logging.set_verbosity(tensorboard_logging.ERROR)
@@ -112,6 +115,7 @@
tensorboard_logging.warn("this should work")
self.assertEqual(1, self.logged_message_count)
+ @test_util.run_v1_only("b/120545219")
def testSummaryWriterFailsAfterClear(self):
tensorboard_logging._clear_summary_writer()
with self.assertRaises(RuntimeError):
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index 51f49ca..8ba6abd 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -53,14 +53,14 @@
self.setUp()
with self.session(use_gpu=use_gpu):
var = variables.VariableV1(x)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(x, self.evaluate(var))
apply_sgd = training_ops.apply_gradient_descent(var, alpha, delta)
out = self.evaluate(apply_sgd)
self.assertShapeEqual(out, apply_sgd)
self.assertAllCloseAccordingToType(x - alpha * delta, out)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testApplyGradientDescent(self):
for (dtype, use_gpu) in itertools.product(
[np.float16, np.float32, np.float64], [False, True]):
@@ -74,7 +74,7 @@
with self.session(use_gpu=use_gpu):
var = variables.VariableV1(x)
accum = variables.VariableV1(y)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(x, self.evaluate(var))
apply_adagrad = training_ops.apply_adagrad(var, accum, lr, grad)
@@ -99,7 +99,7 @@
var = variables.VariableV1(x)
accum = variables.VariableV1(y)
linear = variables.VariableV1(z)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(x, self.evaluate(var))
apply_ftrl = training_ops.apply_ftrl(var, accum, linear, grad, lr, l1, l2,
@@ -129,7 +129,7 @@
self.assertAllClose(linear_update, self.evaluate(linear))
self.assertAllClose(expected_out, out)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testApplyAdagrad(self):
for (dtype, use_gpu) in itertools.product(
[np.float16, np.float32, np.float64], [False, True]):
@@ -139,7 +139,7 @@
grad = np.arange(100).astype(dtype)
self._testTypesForAdagrad(x, y, lr, grad, use_gpu)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testApplyFtrl(self):
for dtype in [np.float16, np.float32, np.float64]:
x = np.arange(100).astype(dtype)
@@ -156,7 +156,7 @@
with self.session(use_gpu=False):
var = variables.VariableV1(x)
accum = variables.VariableV1(y)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(x, self.evaluate(var))
sparse_apply_adagrad = training_ops.sparse_apply_adagrad(
@@ -187,7 +187,7 @@
var = variables.VariableV1(x)
accum = variables.VariableV1(y)
linear = variables.VariableV1(z)
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(x, self.evaluate(var))
sparse_apply_ftrl = training_ops.sparse_apply_ftrl(
@@ -211,7 +211,7 @@
self.assertAllCloseAccordingToType(y[index] + grad[i] * grad[i],
self.evaluate(accum)[index])
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSparseApplyAdagrad(self):
for (dtype, index_type) in itertools.product(
[np.float16, np.float32, np.float64], [np.int32, np.int64]):
@@ -225,7 +225,7 @@
indices = np.array([0, 2]).astype(index_type)
self._testTypesForSparseAdagrad(x, y, lr, grad, indices)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSparseApplyAdagradDim1(self):
for (dtype, index_type) in itertools.product(
[np.float16, np.float32, np.float64], [np.int32, np.int64]):
@@ -239,7 +239,7 @@
indices = np.array([0, 2]).astype(index_type)
self._testTypesForSparseAdagrad(x, y, lr, grad, indices)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testSparseApplyFtrlDim1(self):
for (dtype, index_type) in itertools.product(
[np.float16, np.float32, np.float64], [np.int32, np.int64]):
@@ -255,7 +255,7 @@
indices = np.array([0, 2]).astype(index_type)
self._testTypesForSparseFtrl(x, y, z, lr, grad, indices)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only("b/120545219")
def testApplyAdam(self):
for dtype, use_gpu in itertools.product(
[np.float16, np.float32, np.float64], [False, True]):
@@ -285,7 +285,7 @@
beta2_power_t = variables.VariableV1(beta2_power)
lr_t = constant_op.constant(lr, self._toType(var.dtype), [])
epsilon_t = constant_op.constant(epsilon, self._toType(var.dtype), [])
- variables.global_variables_initializer().run()
+ self.evaluate(variables.global_variables_initializer())
self.assertAllCloseAccordingToType(var, self.evaluate(var_t))
new_var, _, _ = self._adamUpdateNumpy(var, grad, t, m, v, lr, beta1,
diff --git a/tensorflow/python/training/training_util_test.py b/tensorflow/python/training/training_util_test.py
index 3317008..3f9858a 100644
--- a/tensorflow/python/training/training_util_test.py
+++ b/tensorflow/python/training/training_util_test.py
@@ -27,6 +27,7 @@
from tensorflow.python.training import training_util
+@test_util.run_v1_only('b/120545219')
class GlobalStepTest(test.TestCase):
def _assert_global_step(self, global_step, expected_dtype=dtypes.int64):
@@ -47,7 +48,6 @@
self.assertRaisesRegexp(TypeError, 'does not have integer type',
training_util.get_global_step, g)
- @test_util.run_deprecated_v1
def test_invalid_shape(self):
with ops.Graph().as_default() as g:
self.assertIsNone(training_util.get_global_step())
@@ -72,7 +72,6 @@
training_util.create_global_step, g)
self._assert_global_step(training_util.create_global_step(ops.Graph()))
- @test_util.run_deprecated_v1
def test_get_global_step(self):
with ops.Graph().as_default() as g:
self.assertIsNone(training_util.get_global_step())
@@ -93,6 +92,7 @@
self._assert_global_step(training_util.get_or_create_global_step(g))
+@test_util.run_v1_only('b/120545219')
class GlobalStepReadTest(test.TestCase):
def test_global_step_read_is_none_if_there_is_no_global_step(self):
diff --git a/tensorflow/python/training/warm_starting_util.py b/tensorflow/python/training/warm_starting_util.py
index 8c97f10..1382b8c 100644
--- a/tensorflow/python/training/warm_starting_util.py
+++ b/tensorflow/python/training/warm_starting_util.py
@@ -28,7 +28,7 @@
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import checkpoint_ops
from tensorflow.python.training import checkpoint_utils
-from tensorflow.python.training import saver
+from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.util.tf_export import tf_export
@@ -139,7 +139,7 @@
Returns:
Name of the `var`
"""
- name_to_var_dict = saver.BaseSaverBuilder.OpListToDict(var)
+ name_to_var_dict = saveable_object_util.op_list_to_dict(var)
if len(name_to_var_dict) > 1:
raise TypeError("`var` = %s passed as arg violates the constraints. "
"name_to_var_dict = %s" % (var, name_to_var_dict))
diff --git a/tensorflow/python/util/dispatch.py b/tensorflow/python/util/dispatch.py
index e7a56b5..e94e334 100644
--- a/tensorflow/python/util/dispatch.py
+++ b/tensorflow/python/util/dispatch.py
@@ -166,15 +166,14 @@
def add_dispatch_list(target):
"""Decorator that adds a dispatch_list attribute to an op."""
- assert not hasattr(target, DISPATCH_ATTR)
+ if hasattr(target, DISPATCH_ATTR):
+ raise AssertionError("%s already has a dispatch list" % target)
setattr(target, DISPATCH_ATTR, [])
return target
def add_dispatch_support(target):
"""Decorator that adds a dispatch handling wrapper to an op."""
- add_dispatch_list(target)
-
def wrapper(*args, **kwargs):
"""Call target, and fall back on dispatchers if there is a TypeError."""
try:
@@ -188,5 +187,5 @@
else:
raise
- setattr(wrapper, DISPATCH_ATTR, [])
+ add_dispatch_list(wrapper)
return tf_decorator.make_decorator(target, wrapper)
diff --git a/tensorflow/python/util/tf_decorator.py b/tensorflow/python/util/tf_decorator.py
index 0cfc836..f018e1a 100644
--- a/tensorflow/python/util/tf_decorator.py
+++ b/tensorflow/python/util/tf_decorator.py
@@ -98,6 +98,9 @@
if hasattr(target, '__doc__'):
decorator_func.__doc__ = decorator.__doc__
decorator_func.__wrapped__ = target
+ # Keeping a second handle to `target` allows callers to detect whether the
+ # decorator was modified using `rewrap`.
+ decorator_func.__original_wrapped__ = target
return decorator_func
@@ -173,6 +176,8 @@
decorators.append(getattr(cur, '_tf_decorator'))
else:
break
+ if not hasattr(decorators[-1], 'decorated_target'):
+ break
cur = decorators[-1].decorated_target
return decorators, cur
diff --git a/tensorflow/python/util/tf_export.py b/tensorflow/python/util/tf_export.py
index ec70cae..74afc37 100644
--- a/tensorflow/python/util/tf_export.py
+++ b/tensorflow/python/util/tf_export.py
@@ -147,6 +147,94 @@
return None
+def get_v1_names(symbol):
+ """Get a list of TF 1.* names for this symbol.
+
+ Args:
+ symbol: symbol to get API names for.
+
+ Returns:
+ List of all API names for this symbol including TensorFlow and
+ Estimator names.
+ """
+ names_v1 = []
+ tensorflow_api_attr_v1 = API_ATTRS_V1[TENSORFLOW_API_NAME].names
+ estimator_api_attr_v1 = API_ATTRS_V1[ESTIMATOR_API_NAME].names
+
+ if not hasattr(symbol, tensorflow_api_attr_v1):
+ return names_v1
+ if tensorflow_api_attr_v1 in symbol.__dict__:
+ names_v1.extend(getattr(symbol, tensorflow_api_attr_v1))
+ if estimator_api_attr_v1 in symbol.__dict__:
+ names_v1.extend(getattr(symbol, estimator_api_attr_v1))
+ return names_v1
+
+
+def get_v2_names(symbol):
+ """Get a list of TF 2.0 names for this symbol.
+
+ Args:
+ symbol: symbol to get API names for.
+
+ Returns:
+ List of all API names for this symbol including TensorFlow and
+ Estimator names.
+ """
+ names_v2 = []
+ tensorflow_api_attr = API_ATTRS[TENSORFLOW_API_NAME].names
+ estimator_api_attr = API_ATTRS[ESTIMATOR_API_NAME].names
+
+ if not hasattr(symbol, tensorflow_api_attr):
+ return names_v2
+ if tensorflow_api_attr in symbol.__dict__:
+ names_v2.extend(getattr(symbol, tensorflow_api_attr))
+ if estimator_api_attr in symbol.__dict__:
+ names_v2.extend(getattr(symbol, estimator_api_attr))
+ return names_v2
+
+
+def get_v1_constants(module):
+ """Get a list of TF 1.* constants in this module.
+
+ Args:
+ module: TensorFlow module.
+
+ Returns:
+ List of all API constants under the given module including TensorFlow and
+ Estimator constants.
+ """
+ constants_v1 = []
+ tensorflow_constants_attr_v1 = API_ATTRS_V1[TENSORFLOW_API_NAME].constants
+ estimator_constants_attr_v1 = API_ATTRS_V1[ESTIMATOR_API_NAME].constants
+
+ if hasattr(module, tensorflow_constants_attr_v1):
+ constants_v1.extend(getattr(module, tensorflow_constants_attr_v1))
+ if hasattr(module, estimator_constants_attr_v1):
+ constants_v1.extend(getattr(module, estimator_constants_attr_v1))
+ return constants_v1
+
+
+def get_v2_constants(module):
+ """Get a list of TF 2.0 constants in this module.
+
+ Args:
+ module: TensorFlow module.
+
+ Returns:
+ List of all API constants under the given module including TensorFlow and
+ Estimator constants.
+ """
+ constants_v2 = []
+ tensorflow_constants_attr = API_ATTRS[TENSORFLOW_API_NAME].constants
+ estimator_constants_attr = API_ATTRS[ESTIMATOR_API_NAME].constants
+
+ if hasattr(module, tensorflow_constants_attr):
+ constants_v2.extend(getattr(module, tensorflow_constants_attr))
+ if hasattr(module, estimator_constants_attr):
+ constants_v2.extend(getattr(module, estimator_constants_attr))
+ return constants_v2
+
+
class api_export(object): # pylint: disable=invalid-name
"""Provides ways to export symbols to the TensorFlow API."""
diff --git a/tensorflow/python/util/tf_export_test.py b/tensorflow/python/util/tf_export_test.py
index a0fac8b..2062579 100644
--- a/tensorflow/python/util/tf_export_test.py
+++ b/tensorflow/python/util/tf_export_test.py
@@ -62,6 +62,10 @@
del symbol._tf_api_names
if hasattr(symbol, '_tf_api_names_v1'):
del symbol._tf_api_names_v1
+ if hasattr(symbol, '_estimator_api_names'):
+ del symbol._estimator_api_names
+ if hasattr(symbol, '_estimator_api_names_v1'):
+ del symbol._estimator_api_names_v1
def _CreateMockModule(self, name):
mock_module = self.MockModule(name)
@@ -74,6 +78,10 @@
decorated_function = export_decorator(_test_function)
self.assertEquals(decorated_function, _test_function)
self.assertEquals(('nameA', 'nameB'), decorated_function._tf_api_names)
+ self.assertEquals(['nameA', 'nameB'],
+ tf_export.get_v1_names(decorated_function))
+ self.assertEquals(['nameA', 'nameB'],
+ tf_export.get_v2_names(decorated_function))
def testExportMultipleFunctions(self):
export_decorator1 = tf_export.tf_export('nameA', 'nameB')
@@ -95,6 +103,22 @@
export_decorator_b(TestClassB)
self.assertEquals(('TestClassA1',), TestClassA._tf_api_names)
self.assertEquals(('TestClassB1',), TestClassB._tf_api_names)
+ self.assertEquals(['TestClassA1'], tf_export.get_v1_names(TestClassA))
+ self.assertEquals(['TestClassB1'], tf_export.get_v1_names(TestClassB))
+
+ def testExportClassInEstimator(self):
+ export_decorator_a = tf_export.tf_export('TestClassA1')
+ export_decorator_a(TestClassA)
+ self.assertEquals(('TestClassA1',), TestClassA._tf_api_names)
+
+ export_decorator_b = tf_export.estimator_export(
+ 'estimator.TestClassB1')
+ export_decorator_b(TestClassB)
+ self.assertTrue('_tf_api_names' not in TestClassB.__dict__)
+ self.assertEquals(('TestClassA1',), TestClassA._tf_api_names)
+ self.assertEquals(['TestClassA1'], tf_export.get_v1_names(TestClassA))
+ self.assertEquals(['estimator.TestClassB1'],
+ tf_export.get_v1_names(TestClassB))
def testExportSingleConstant(self):
module1 = self._CreateMockModule('module1')
@@ -103,6 +127,10 @@
export_decorator.export_constant('module1', 'test_constant')
self.assertEquals([(('NAME_A', 'NAME_B'), 'test_constant')],
module1._tf_api_constants)
+ self.assertEquals([(('NAME_A', 'NAME_B'), 'test_constant')],
+ tf_export.get_v1_constants(module1))
+ self.assertEquals([(('NAME_A', 'NAME_B'), 'test_constant')],
+ tf_export.get_v2_constants(module1))
def testExportMultipleConstants(self):
module1 = self._CreateMockModule('module1')
diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py
index ca6710b..63de4a7 100644
--- a/tensorflow/python/util/tf_should_use.py
+++ b/tensorflow/python/util/tf_should_use.py
@@ -23,6 +23,7 @@
import six # pylint: disable=unused-import
+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
@@ -32,7 +33,8 @@
"""Object stored in TFShouldUse-wrapped objects.
When it is deleted it will emit a warning or error if its `sate` method
- has not been called by time of deletion.
+ has not been called by time of deletion, and Tensorflow is not executing
+ eagerly outside of functions.
"""
def __init__(self, type_, repr_, stack_frame, fatal_error_if_unsated):
@@ -50,6 +52,8 @@
self._logging_module = None
def __del__(self):
+ if ops.executing_eagerly_outside_functions():
+ return
if self._sated:
return
if self._fatal_error_if_unsated:
diff --git a/tensorflow/stream_executor/cuda/cuda_helpers.h b/tensorflow/stream_executor/cuda/cuda_helpers.h
index d55706c..dc0dc69 100644
--- a/tensorflow/stream_executor/cuda/cuda_helpers.h
+++ b/tensorflow/stream_executor/cuda/cuda_helpers.h
@@ -25,7 +25,6 @@
#include <complex>
#include "cuda/include/cuComplex.h"
-#include "cuda/include/cuda.h"
namespace stream_executor {
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.-ragged-tensor.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-ragged-tensor.pbtxt
new file mode 100644
index 0000000..c0ed956
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-ragged-tensor.pbtxt
@@ -0,0 +1,125 @@
+path: "tensorflow.RaggedTensor"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "flat_values"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "nested_row_splits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "ragged_rank"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "row_splits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "values"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'values\', \'row_splits\', \'cached_row_lengths\', \'cached_value_rowids\', \'cached_nrows\', \'internal\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'False\'], "
+ }
+ member_method {
+ name: "bounding_shape"
+ argspec: "args=[\'self\', \'axis\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_nested_row_lengths"
+ argspec: "args=[\'cls\', \'flat_values\', \'nested_row_lengths\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_nested_row_splits"
+ argspec: "args=[\'cls\', \'flat_values\', \'nested_row_splits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_nested_value_rowids"
+ argspec: "args=[\'cls\', \'flat_values\', \'nested_value_rowids\', \'nested_nrows\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_row_lengths"
+ argspec: "args=[\'cls\', \'values\', \'row_lengths\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_row_limits"
+ argspec: "args=[\'cls\', \'values\', \'row_limits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_row_splits"
+ argspec: "args=[\'cls\', \'values\', \'row_splits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_row_starts"
+ argspec: "args=[\'cls\', \'values\', \'row_starts\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_sparse"
+ argspec: "args=[\'cls\', \'st_input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_tensor"
+ argspec: "args=[\'cls\', \'tensor\', \'lengths\', \'padding\', \'ragged_rank\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'1\', \'None\'], "
+ }
+ member_method {
+ name: "from_value_rowids"
+ argspec: "args=[\'cls\', \'values\', \'value_rowids\', \'nrows\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "nested_row_lengths"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "nrows"
+ argspec: "args=[\'self\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'int64\'>\", \'None\'], "
+ }
+ member_method {
+ name: "row_lengths"
+ argspec: "args=[\'self\', \'axis\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "row_limits"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "row_starts"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "to_list"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "to_sparse"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "to_tensor"
+ argspec: "args=[\'self\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "value_rowids"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "with_flat_values"
+ argspec: "args=[\'self\', \'new_values\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "with_values"
+ argspec: "args=[\'self\', \'new_values\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt
index af7fc9d..62d8ea9 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.-variable.pbtxt
@@ -65,6 +65,10 @@
argspec: "args=[\'self\', \'delta\', \'use_locking\', \'name\', \'read_value\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'True\'], "
}
member_method {
+ name: "batch_scatter_update"
+ argspec: "args=[\'self\', \'sparse_delta\', \'use_locking\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
name: "count_up_to"
argspec: "args=[\'self\', \'limit\'], varargs=None, keywords=None, defaults=None"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.-feature.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.-feature.pbtxt
new file mode 100644
index 0000000..a71da11
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.-feature.pbtxt
@@ -0,0 +1,28 @@
+path: "tensorflow.autograph.experimental.Feature"
+tf_class {
+ is_instance: "<enum \'Feature\'>"
+ member {
+ name: "ALL"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "AUTO_CONTROL_DEPS"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "DECORATORS"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "ERROR_REWRITING"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "LISTS"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "NAME_SCOPES"
+ mtype: "<enum \'Feature\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.-verbosity.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.-verbosity.pbtxt
new file mode 100644
index 0000000..c4d5b77
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.-verbosity.pbtxt
@@ -0,0 +1,12 @@
+path: "tensorflow.autograph.experimental.Verbosity"
+tf_class {
+ is_instance: "<enum \'Verbosity\'>"
+ member {
+ name: "BRIEF"
+ mtype: "<enum \'Verbosity\'>"
+ }
+ member {
+ name: "VERBOSE"
+ mtype: "<enum \'Verbosity\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.pbtxt
new file mode 100644
index 0000000..5747dac
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.autograph.experimental.pbtxt
@@ -0,0 +1,11 @@
+path: "tensorflow.autograph.experimental"
+tf_module {
+ member {
+ name: "Feature"
+ mtype: "<class \'enum.EnumMeta\'>"
+ }
+ member {
+ name: "Verbosity"
+ mtype: "<class \'enum.EnumMeta\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.autograph.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.autograph.pbtxt
new file mode 100644
index 0000000..34bdab9
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.autograph.pbtxt
@@ -0,0 +1,15 @@
+path: "tensorflow.autograph"
+tf_module {
+ member {
+ name: "experimental"
+ mtype: "<type \'module\'>"
+ }
+ member_method {
+ name: "to_code"
+ argspec: "args=[\'entity\', \'recursive\', \'arg_values\', \'arg_types\', \'indentation\', \'optional_features\', \'experimental_partial_types\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\', \' \', \'Feature.ALL\', \'None\'], "
+ }
+ member_method {
+ name: "to_graph"
+ argspec: "args=[\'entity\', \'recursive\', \'verbose\', \'arg_values\', \'arg_types\', \'strip_decorators\', \'optional_features\', \'experimental_partial_types\'], varargs=None, keywords=None, defaults=[\'True\', \'Verbosity.VERBOSE\', \'None\', \'None\', \'None\', \'Feature.ALL\', \'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt
index f59082b..f7d388d 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.-dataset.pbtxt
@@ -5,15 +5,15 @@
is_instance: "<type \'object\'>"
member {
name: "output_classes"
- mtype: "<class \'abc.abstractproperty\'>"
+ mtype: "<type \'property\'>"
}
member {
name: "output_shapes"
- mtype: "<class \'abc.abstractproperty\'>"
+ mtype: "<type \'property\'>"
}
member {
name: "output_types"
- mtype: "<class \'abc.abstractproperty\'>"
+ mtype: "<type \'property\'>"
}
member_method {
name: "__init__"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-dataset-structure.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-dataset-structure.pbtxt
new file mode 100644
index 0000000..dcb304f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-dataset-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.DatasetStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'element_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-nested-structure.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-nested-structure.pbtxt
new file mode 100644
index 0000000..b4b066e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-nested-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.NestedStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.NestedStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'nested_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt
index 9ca7582..3b7ad64 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optimization-options.pbtxt
@@ -4,6 +4,10 @@
is_instance: "<class \'tensorflow.python.data.util.options.OptionsBase\'>"
is_instance: "<type \'object\'>"
member {
+ name: "apply_default_optimizations"
+ mtype: "<type \'property\'>"
+ }
+ member {
name: "filter_fusion"
mtype: "<type \'property\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optional-structure.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optional-structure.pbtxt
new file mode 100644
index 0000000..bf41c1d
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-optional-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.OptionalStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.ops.optional_ops.OptionalStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'value_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sparse-tensor-structure.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sparse-tensor-structure.pbtxt
new file mode 100644
index 0000000..f97376b
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-sparse-tensor-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.SparseTensorStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.SparseTensorStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtype\', \'dense_shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-structure.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-structure.pbtxt
new file mode 100644
index 0000000..a99db45
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-structure.pbtxt
@@ -0,0 +1,16 @@
+path: "tensorflow.data.experimental.Structure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-tensor-structure.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-tensor-structure.pbtxt
new file mode 100644
index 0000000..f5c8864
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.-tensor-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.TensorStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.TensorStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtype\', \'shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt
index f981b1a..2d11590 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.data.experimental.pbtxt
@@ -13,10 +13,18 @@
mtype: "<type \'type\'>"
}
member {
+ name: "DatasetStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "INFINITE_CARDINALITY"
mtype: "<type \'int\'>"
}
member {
+ name: "NestedStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "OptimizationOptions"
mtype: "<type \'type\'>"
}
@@ -25,6 +33,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "OptionalStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "RandomDataset"
mtype: "<type \'type\'>"
}
@@ -33,6 +45,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "SparseTensorStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "SqlDataset"
mtype: "<type \'type\'>"
}
@@ -45,10 +61,18 @@
mtype: "<type \'type\'>"
}
member {
+ name: "Structure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "TFRecordWriter"
mtype: "<type \'type\'>"
}
member {
+ name: "TensorStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "ThreadingOptions"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
new file mode 100644
index 0000000..81224f0
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
@@ -0,0 +1,142 @@
+path: "tensorflow.distribute.MirroredStrategy"
+tf_class {
+ is_instance: "<class \'tensorflow.python.distribute.mirrored_strategy.MirroredStrategy\'>"
+ is_instance: "<class \'tensorflow.python.distribute.distribute_lib.DistributionStrategy\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "between_graph"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "extended"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "num_replicas_in_sync"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "parameter_devices"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "require_static_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "should_checkpoint"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "should_init"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "should_save_summary"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "worker_devices"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'devices\', \'cross_device_ops\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "batch_reduce"
+ argspec: "args=[\'self\', \'aggregation\', \'value_destination_pairs\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "broadcast"
+ argspec: "args=[\'self\', \'tensor\', \'destinations\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "call_for_each_replica"
+ argspec: "args=[\'self\', \'fn\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "colocate_vars_with"
+ argspec: "args=[\'self\', \'colocate_with_variable\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "configure"
+ argspec: "args=[\'self\', \'session_config\', \'cluster_spec\', \'task_type\', \'task_id\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "distribute_dataset"
+ argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "experimental_finalize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "experimental_initialize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "experimental_run"
+ argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "finalize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "group"
+ argspec: "args=[\'self\', \'value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "initialize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "make_dataset_iterator"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "make_input_fn_iterator"
+ argspec: "args=[\'self\', \'input_fn\', \'replication_mode\'], varargs=None, keywords=None, defaults=[\'InputReplicationMode.PER_WORKER\'], "
+ }
+ member_method {
+ name: "non_slot_devices"
+ argspec: "args=[\'self\', \'var_list\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "read_var"
+ argspec: "args=[\'self\', \'v\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "run_steps_on_dataset"
+ argspec: "args=[\'self\', \'fn\', \'iterator\', \'iterations\', \'initial_loop_values\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "scope"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "unwrap"
+ argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "update"
+ argspec: "args=[\'self\', \'var\', \'fn\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "update_config_proto"
+ argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "update_non_slot"
+ argspec: "args=[\'self\', \'colocate_with\', \'fn\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "value_container"
+ argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-server.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-server.pbtxt
similarity index 95%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-server.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.distribute.-server.pbtxt
index 9b8f185..6c39bf4 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-server.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-server.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.Server"
+path: "tensorflow.distribute.Server"
tf_class {
is_instance: "<class \'tensorflow.python.training.server_lib.Server\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
index 9eb73d2..63b6584 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
@@ -75,6 +75,10 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "experimental_run"
+ argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "finalize"
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt
index 4d833b5..31dc6e0 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.pbtxt
@@ -9,6 +9,10 @@
mtype: "<class \'enum.EnumMeta\'>"
}
member {
+ name: "MirroredStrategy"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "ReduceOp"
mtype: "<class \'enum.EnumMeta\'>"
}
@@ -17,6 +21,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "Server"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "Strategy"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt
index 848fc30..01b870a8 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.dtypes.pbtxt
@@ -45,6 +45,10 @@
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
member {
+ name: "int16"
+ mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
+ }
+ member {
name: "int32"
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-checkpoint-saver-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-checkpoint-saver-hook.pbtxt
similarity index 95%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-checkpoint-saver-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-checkpoint-saver-hook.pbtxt
index c3037ba..f9e1504 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-checkpoint-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-checkpoint-saver-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.CheckpointSaverHook"
+path: "tensorflow.estimator.CheckpointSaverHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.CheckpointSaverHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.-checkpoint-saver-listener.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-checkpoint-saver-listener.pbtxt
new file mode 100644
index 0000000..111b758
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-checkpoint-saver-listener.pbtxt
@@ -0,0 +1,24 @@
+path: "tensorflow.estimator.CheckpointSaverListener"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.CheckpointSaverListener\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "after_save"
+ argspec: "args=[\'self\', \'session\', \'global_step_value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "before_save"
+ argspec: "args=[\'self\', \'session\', \'global_step_value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "begin"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "end"
+ argspec: "args=[\'self\', \'session\', \'global_step_value\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-feed-fn-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-feed-fn-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-feed-fn-hook.pbtxt
copy to tensorflow/tools/api/golden/v1/tensorflow.estimator.-feed-fn-hook.pbtxt
index 7bec4d0..f24de49 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-feed-fn-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-feed-fn-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.FeedFnHook"
+path: "tensorflow.estimator.FeedFnHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.FeedFnHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-final-ops-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-final-ops-hook.pbtxt
similarity index 96%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-final-ops-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-final-ops-hook.pbtxt
index 31cf9aa..6651170 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-final-ops-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-final-ops-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.FinalOpsHook"
+path: "tensorflow.estimator.FinalOpsHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.FinalOpsHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-global-step-waiter-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-global-step-waiter-hook.pbtxt
similarity index 95%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-global-step-waiter-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-global-step-waiter-hook.pbtxt
index 1474486..37db48b 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-global-step-waiter-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-global-step-waiter-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.GlobalStepWaiterHook"
+path: "tensorflow.estimator.GlobalStepWaiterHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.GlobalStepWaiterHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-logging-tensor-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-logging-tensor-hook.pbtxt
similarity index 95%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-logging-tensor-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-logging-tensor-hook.pbtxt
index 9801c05..425f016 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-logging-tensor-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-logging-tensor-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.LoggingTensorHook"
+path: "tensorflow.estimator.LoggingTensorHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.LoggingTensorHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.-nan-loss-during-training-error.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-nan-loss-during-training-error.pbtxt
new file mode 100644
index 0000000..6cf6e17
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-nan-loss-during-training-error.pbtxt
@@ -0,0 +1,12 @@
+path: "tensorflow.estimator.NanLossDuringTrainingError"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.NanLossDuringTrainingError\'>"
+ is_instance: "<type \'exceptions.RuntimeError\'>"
+ member {
+ name: "args"
+ mtype: "<type \'getset_descriptor\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-nan-tensor-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-nan-tensor-hook.pbtxt
similarity index 95%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-nan-tensor-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-nan-tensor-hook.pbtxt
index 7d1c89f..82293c2 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-nan-tensor-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-nan-tensor-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.NanTensorHook"
+path: "tensorflow.estimator.NanTensorHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.NanTensorHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-profiler-hook.pbtxt
similarity index 77%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt
copy to tensorflow/tools/api/golden/v1/tensorflow.estimator.-profiler-hook.pbtxt
index 697c366..65b5fb1 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-profiler-hook.pbtxt
@@ -1,11 +1,11 @@
-path: "tensorflow.train.SummarySaverHook"
+path: "tensorflow.estimator.ProfilerHook"
tf_class {
- is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.SummarySaverHook\'>"
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.ProfilerHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
is_instance: "<type \'object\'>"
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'save_steps\', \'save_secs\', \'output_dir\', \'summary_writer\', \'scaffold\', \'summary_op\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
+ argspec: "args=[\'self\', \'save_steps\', \'save_secs\', \'output_dir\', \'show_dataflow\', \'show_memory\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'\', \'True\', \'False\'], "
}
member_method {
name: "after_create_session"
@@ -25,6 +25,6 @@
}
member_method {
name: "end"
- argspec: "args=[\'self\', \'session\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ argspec: "args=[\'self\', \'session\'], varargs=None, keywords=None, defaults=None"
}
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.-second-or-step-timer.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-second-or-step-timer.pbtxt
new file mode 100644
index 0000000..64051d2
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-second-or-step-timer.pbtxt
@@ -0,0 +1,26 @@
+path: "tensorflow.estimator.SecondOrStepTimer"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.SecondOrStepTimer\'>"
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks._HookTimer\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'every_secs\', \'every_steps\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "last_triggered_step"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reset"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "should_trigger_for_step"
+ argspec: "args=[\'self\', \'step\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "update_last_triggered_step"
+ argspec: "args=[\'self\', \'step\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-args.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-args.pbtxt
new file mode 100644
index 0000000..b375c74
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-args.pbtxt
@@ -0,0 +1,27 @@
+path: "tensorflow.estimator.SessionRunArgs"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunArgs\'>"
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunArgs\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "feed_dict"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "fetches"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "options"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-context.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-context.pbtxt
new file mode 100644
index 0000000..cb4ac9f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-context.pbtxt
@@ -0,0 +1,25 @@
+path: "tensorflow.estimator.SessionRunContext"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunContext\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "original_args"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "session"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "stop_requested"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'original_args\', \'session\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "request_stop"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-session-run-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-hook.pbtxt
similarity index 94%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-session-run-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-hook.pbtxt
index db1aa24..54e9ad9 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-session-run-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.SessionRunHook"
+path: "tensorflow.estimator.SessionRunHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-values.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-values.pbtxt
new file mode 100644
index 0000000..6788141
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-session-run-values.pbtxt
@@ -0,0 +1,27 @@
+path: "tensorflow.estimator.SessionRunValues"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunValues\'>"
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunValues\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "options"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "results"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "run_metadata"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-step-counter-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-step-counter-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-step-counter-hook.pbtxt
copy to tensorflow/tools/api/golden/v1/tensorflow.estimator.-step-counter-hook.pbtxt
index 13261f6..4368e04 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-step-counter-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-step-counter-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.StepCounterHook"
+path: "tensorflow.estimator.StepCounterHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.StepCounterHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-stop-at-step-hook.pbtxt
similarity index 95%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.estimator.-stop-at-step-hook.pbtxt
index e388599..938b189 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-stop-at-step-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.StopAtStepHook"
+path: "tensorflow.estimator.StopAtStepHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.StopAtStepHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-summary-saver-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt
copy to tensorflow/tools/api/golden/v1/tensorflow.estimator.-summary-saver-hook.pbtxt
index 697c366..1041573 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.-summary-saver-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.SummarySaverHook"
+path: "tensorflow.estimator.SummarySaverHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.SummarySaverHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt
index aba1202..5a2a01c 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt
@@ -1,6 +1,6 @@
path: "tensorflow.estimator.experimental.InMemoryEvaluatorHook"
tf_class {
- is_instance: "<class \'tensorflow_estimator.python.estimator.hooks.InMemoryEvaluatorHook\'>"
+ is_instance: "<class \'tensorflow_estimator.python.estimator.hooks.hooks.InMemoryEvaluatorHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
is_instance: "<type \'object\'>"
member_method {
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.estimator.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.estimator.pbtxt
index c5b0085..6f57505 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.estimator.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.estimator.pbtxt
@@ -25,6 +25,14 @@
mtype: "<type \'type\'>"
}
member {
+ name: "CheckpointSaverHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "CheckpointSaverListener"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "DNNClassifier"
mtype: "<type \'type\'>"
}
@@ -65,10 +73,22 @@
mtype: "<type \'type\'>"
}
member {
+ name: "FeedFnHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "FinalExporter"
mtype: "<type \'type\'>"
}
member {
+ name: "FinalOpsHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "GlobalStepWaiterHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "LatestExporter"
mtype: "<type \'type\'>"
}
@@ -85,14 +105,62 @@
mtype: "<type \'type\'>"
}
member {
+ name: "LoggingTensorHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "ModeKeys"
mtype: "<type \'type\'>"
}
member {
+ name: "NanLossDuringTrainingError"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "NanTensorHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "ProfilerHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "RunConfig"
mtype: "<type \'type\'>"
}
member {
+ name: "SecondOrStepTimer"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunArgs"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunContext"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunValues"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "StepCounterHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "StopAtStepHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SummarySaverHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "TrainSpec"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.io.gfile.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.io.gfile.pbtxt
index 93d9b0f..cfa3372 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.io.gfile.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.io.gfile.pbtxt
@@ -46,6 +46,6 @@
}
member_method {
name: "walk"
- argspec: "args=[\'top\', \'topdown\', \'onerror\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ argspec: "args=[\'top\', \'topdown\', \'onerror\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
}
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.-sensitivity-at-specificity.pbtxt
similarity index 92%
copy from tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
copy to tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.-sensitivity-at-specificity.pbtxt
index 33226a2..aa77d19 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.-sensitivity-at-specificity.pbtxt
@@ -1,7 +1,7 @@
-path: "tensorflow.metrics.FalseNegatives"
+path: "tensorflow.keras.metrics.SensitivityAtSpecificity"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalseNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SensitivityAtSpecificity\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SensitivitySpecificityBase\'>"
is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
@@ -84,7 +84,7 @@
}
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ argspec: "args=[\'self\', \'specificity\', \'num_thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'200\', \'None\', \'None\'], "
}
member_method {
name: "add_loss"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.-specificity-at-sensitivity.pbtxt
similarity index 92%
rename from tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.-specificity-at-sensitivity.pbtxt
index 33226a2..67857aa 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.-specificity-at-sensitivity.pbtxt
@@ -1,7 +1,7 @@
-path: "tensorflow.metrics.FalseNegatives"
+path: "tensorflow.keras.metrics.SpecificityAtSensitivity"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalseNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SpecificityAtSensitivity\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SensitivitySpecificityBase\'>"
is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
@@ -84,7 +84,7 @@
}
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ argspec: "args=[\'self\', \'sensitivity\', \'num_thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'200\', \'None\', \'None\'], "
}
member_method {
name: "add_loss"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.pbtxt
index 8cab17e..905021d 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.metrics.pbtxt
@@ -33,10 +33,18 @@
mtype: "<type \'type\'>"
}
member {
+ name: "SensitivityAtSpecificity"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "SparseCategoricalAccuracy"
mtype: "<type \'type\'>"
}
member {
+ name: "SpecificityAtSensitivity"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "TrueNegatives"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adadelta.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adadelta.pbtxt
index b9ce154..5426269 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adadelta.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adadelta.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.Adadelta"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adadelta\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adadelta.Adadelta\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'rho\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'1.0\', \'0.95\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'rho\', \'epsilon\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.95\', \'1e-07\', \'Adadelta\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adagrad.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adagrad.pbtxt
index d0dc9e3..c39fe6b 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adagrad.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adagrad.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.Adagrad"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adagrad\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adagrad.Adagrad\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'0.01\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'initial_accumulator_value\', \'epsilon\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.1\', \'1e-07\', \'Adagrad\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adam.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adam.pbtxt
index 06815fa..05d46d3 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adam.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adam.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.Adam"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adam\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adam.Adam\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'beta_1\', \'beta_2\', \'epsilon\', \'decay\', \'amsgrad\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.999\', \'None\', \'0.0\', \'False\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'beta_1\', \'beta_2\', \'epsilon\', \'amsgrad\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.999\', \'1e-07\', \'False\', \'Adam\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adamax.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adamax.pbtxt
index 47b55fd..78829de 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adamax.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-adamax.pbtxt
@@ -1,15 +1,37 @@
path: "tensorflow.keras.optimizers.Adamax"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adamax\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adamax.Adamax\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adam.Adam\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'beta_1\', \'beta_2\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'0.002\', \'0.9\', \'0.999\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'beta_1\', \'beta_2\', \'epsilon\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.999\', \'1e-07\', \'Adamax\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +42,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +58,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-optimizer.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-optimizer.pbtxt
index 53d64da..58b7f27 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-optimizer.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-optimizer.pbtxt
@@ -1,14 +1,35 @@
path: "tensorflow.keras.optimizers.Optimizer"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -19,6 +40,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -27,7 +56,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-r-m-sprop.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-r-m-sprop.pbtxt
index a1e9b8c..8de796e 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-r-m-sprop.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-r-m-sprop.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.RMSprop"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.RMSprop\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.rmsprop.RMSprop\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'rho\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'rho\', \'momentum\', \'epsilon\', \'centered\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.0\', \'1e-07\', \'False\', \'RMSprop\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-s-g-d.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-s-g-d.pbtxt
index a67fefb..393eeb3 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-s-g-d.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.optimizers.-s-g-d.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.SGD"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.SGD\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.gradient_descent.SGD\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'momentum\', \'decay\', \'nesterov\'], varargs=None, keywords=kwargs, defaults=[\'0.01\', \'0.0\', \'0.0\', \'False\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'momentum\', \'nesterov\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.0\', \'False\', \'SGD\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt
index 1a4098d..9f7b422 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.linalg.pbtxt
@@ -133,6 +133,10 @@
argspec: "args=[\'matrix\', \'rhs\', \'l2_regularizer\', \'fast\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'True\', \'None\'], "
}
member_method {
+ name: "lu"
+ argspec: "args=[\'input\', \'output_idx_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'int32\'>\", \'None\'], "
+ }
+ member_method {
name: "matmul"
argspec: "args=[\'a\', \'b\', \'transpose_a\', \'transpose_b\', \'adjoint_a\', \'adjoint_b\', \'a_is_sparse\', \'b_is_sparse\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'False\', \'False\', \'False\', \'False\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt
new file mode 100644
index 0000000..68c651a
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt
@@ -0,0 +1,12 @@
+path: "tensorflow.lite.OpsSet"
+tf_class {
+ is_instance: "<enum \'OpsSet\'>"
+ member {
+ name: "SELECT_TF_OPS"
+ mtype: "<enum \'OpsSet\'>"
+ }
+ member {
+ name: "TFLITE_BUILTINS"
+ mtype: "<enum \'OpsSet\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.lite.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.lite.pbtxt
index f5013c2..154dd00 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.lite.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.lite.pbtxt
@@ -9,6 +9,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "OpsSet"
+ mtype: "<class \'enum.EnumMeta\'>"
+ }
+ member {
name: "TFLiteConverter"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-accuracy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-accuracy.pbtxt
deleted file mode 100644
index f8e12f8..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.Accuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Accuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'accuracy\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-binary-accuracy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-binary-accuracy.pbtxt
deleted file mode 100644
index b9bc6a7..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-binary-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.BinaryAccuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.BinaryAccuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\', \'threshold\'], varargs=None, keywords=None, defaults=[\'binary_accuracy\', \'None\', \'0.5\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-categorical-accuracy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-categorical-accuracy.pbtxt
deleted file mode 100644
index 0ef75d8..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-categorical-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.CategoricalAccuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.CategoricalAccuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'categorical_accuracy\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-positives.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-positives.pbtxt
deleted file mode 100644
index 9953162..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-positives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.FalsePositives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalsePositives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-mean.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-mean.pbtxt
deleted file mode 100644
index 7fe6d6f..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-mean.pbtxt
+++ /dev/null
@@ -1,192 +0,0 @@
-path: "tensorflow.metrics.Mean"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'mean\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'values\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-precision.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-precision.pbtxt
deleted file mode 100644
index 8c3271a..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-precision.pbtxt
+++ /dev/null
@@ -1,192 +0,0 @@
-path: "tensorflow.metrics.Precision"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Precision\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-recall.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-recall.pbtxt
deleted file mode 100644
index 840a68b..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-recall.pbtxt
+++ /dev/null
@@ -1,192 +0,0 @@
-path: "tensorflow.metrics.Recall"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Recall\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-sparse-categorical-accuracy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-sparse-categorical-accuracy.pbtxt
deleted file mode 100644
index 7bce43f..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-sparse-categorical-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.SparseCategoricalAccuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.SparseCategoricalAccuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'sparse_categorical_accuracy\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-true-negatives.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-true-negatives.pbtxt
deleted file mode 100644
index 83cd5b7..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-true-negatives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.TrueNegatives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.TrueNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-true-positives.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.-true-positives.pbtxt
deleted file mode 100644
index 5b2502e..0000000
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-true-positives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.TruePositives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.TruePositives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.metrics.pbtxt
index f5c267a..e9b996c 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.metrics.pbtxt
@@ -1,49 +1,5 @@
path: "tensorflow.metrics"
tf_module {
- member {
- name: "Accuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "BinaryAccuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "CategoricalAccuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "FalseNegatives"
- mtype: "<type \'type\'>"
- }
- member {
- name: "FalsePositives"
- mtype: "<type \'type\'>"
- }
- member {
- name: "Mean"
- mtype: "<type \'type\'>"
- }
- member {
- name: "Precision"
- mtype: "<type \'type\'>"
- }
- member {
- name: "Recall"
- mtype: "<type \'type\'>"
- }
- member {
- name: "SparseCategoricalAccuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "TrueNegatives"
- mtype: "<type \'type\'>"
- }
- member {
- name: "TruePositives"
- mtype: "<type \'type\'>"
- }
member_method {
name: "accuracy"
argspec: "args=[\'labels\', \'predictions\', \'weights\', \'metrics_collections\', \'updates_collections\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.pbtxt
index 584c74f..4ed4dee 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.pbtxt
@@ -173,6 +173,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "RaggedTensor"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "RandomShuffleQueue"
mtype: "<type \'type\'>"
}
@@ -289,6 +293,10 @@
mtype: "<type \'module\'>"
}
member {
+ name: "autograph"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "bfloat16"
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
@@ -505,6 +513,10 @@
mtype: "<type \'module\'>"
}
member {
+ name: "queue"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "quint16"
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
@@ -513,6 +525,10 @@
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
member {
+ name: "ragged"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "random"
mtype: "<type \'module\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-f-i-f-o-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.queue.-f-i-f-o-queue.pbtxt
similarity index 97%
copy from tensorflow/tools/api/golden/v2/tensorflow.-f-i-f-o-queue.pbtxt
copy to tensorflow/tools/api/golden/v1/tensorflow.queue.-f-i-f-o-queue.pbtxt
index a095616..724ab5f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.-f-i-f-o-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.queue.-f-i-f-o-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.FIFOQueue"
+path: "tensorflow.queue.FIFOQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.FIFOQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.queue.-padding-f-i-f-o-queue.pbtxt
similarity index 97%
rename from tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.queue.-padding-f-i-f-o-queue.pbtxt
index 85306fd..9ef0a4d 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.queue.-padding-f-i-f-o-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.PaddingFIFOQueue"
+path: "tensorflow.queue.PaddingFIFOQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PaddingFIFOQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.queue.-priority-queue.pbtxt
similarity index 97%
rename from tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.queue.-priority-queue.pbtxt
index 02d8037..bb66beb 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.queue.-priority-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.PriorityQueue"
+path: "tensorflow.queue.PriorityQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PriorityQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.queue.-queue-base.pbtxt
similarity index 97%
rename from tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.queue.-queue-base.pbtxt
index a30481a..8faaad2 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.queue.-queue-base.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.QueueBase"
+path: "tensorflow.queue.QueueBase"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.queue.-random-shuffle-queue.pbtxt
similarity index 97%
rename from tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt
rename to tensorflow/tools/api/golden/v1/tensorflow.queue.-random-shuffle-queue.pbtxt
index 82cbf98..31cd503 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.queue.-random-shuffle-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.RandomShuffleQueue"
+path: "tensorflow.queue.RandomShuffleQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.RandomShuffleQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.queue.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.queue.pbtxt
new file mode 100644
index 0000000..c16e95e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.queue.pbtxt
@@ -0,0 +1,23 @@
+path: "tensorflow.queue"
+tf_module {
+ member {
+ name: "FIFOQueue"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "PaddingFIFOQueue"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "PriorityQueue"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "QueueBase"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "RandomShuffleQueue"
+ mtype: "<type \'type\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.ragged.-ragged-tensor-value.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.ragged.-ragged-tensor-value.pbtxt
new file mode 100644
index 0000000..96c895e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.ragged.-ragged-tensor-value.pbtxt
@@ -0,0 +1,41 @@
+path: "tensorflow.ragged.RaggedTensorValue"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.ragged.ragged_tensor_value.RaggedTensorValue\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "flat_values"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "nested_row_splits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "ragged_rank"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "row_splits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "values"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'values\', \'row_splits\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "to_list"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.ragged.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.ragged.pbtxt
new file mode 100644
index 0000000..22ca7e9
--- /dev/null
+++ b/tensorflow/tools/api/golden/v1/tensorflow.ragged.pbtxt
@@ -0,0 +1,31 @@
+path: "tensorflow.ragged"
+tf_module {
+ member {
+ name: "RaggedTensorValue"
+ mtype: "<type \'type\'>"
+ }
+ member_method {
+ name: "constant"
+ argspec: "args=[\'pylist\', \'dtype\', \'ragged_rank\', \'inner_shape\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "constant_value"
+ argspec: "args=[\'pylist\', \'dtype\', \'ragged_rank\', \'inner_shape\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "map_flat_values"
+ argspec: "args=[\'op\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[\'starts\', \'limits\', \'deltas\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "row_splits_to_segment_ids"
+ argspec: "args=[\'splits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "segment_ids_to_row_splits"
+ argspec: "args=[\'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
index 107534e..1eefb1c 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
@@ -1,10 +1,18 @@
path: "tensorflow.random"
tf_module {
member_method {
+ name: "all_candidate_sampler"
+ argspec: "args=[\'true_classes\', \'num_true\', \'num_sampled\', \'unique\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "categorical"
argspec: "args=[\'logits\', \'num_samples\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
}
member_method {
+ name: "fixed_unigram_candidate_sampler"
+ argspec: "args=[\'true_classes\', \'num_true\', \'num_sampled\', \'unique\', \'range_max\', \'vocab_file\', \'distortion\', \'num_reserved_ids\', \'num_shards\', \'shard\', \'unigrams\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'1.0\', \'0\', \'1\', \'0\', \'()\', \'None\', \'None\'], "
+ }
+ member_method {
name: "gamma"
argspec: "args=[\'shape\', \'alpha\', \'beta\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
}
@@ -13,6 +21,10 @@
argspec: "args=[\'op_seed\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "learned_unigram_candidate_sampler"
+ argspec: "args=[\'true_classes\', \'num_true\', \'num_sampled\', \'unique\', \'range_max\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "log_uniform_candidate_sampler"
argspec: "args=[\'true_classes\', \'num_true\', \'num_sampled\', \'unique\', \'range_max\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.strings.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.strings.pbtxt
index a1cd581..ad26ded 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.strings.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.strings.pbtxt
@@ -53,6 +53,14 @@
argspec: "args=[\'string_tensor\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'float32\'>\", \'None\'], "
}
member_method {
+ name: "unicode_decode"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'replace_control_characters\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "unicode_decode_with_offsets"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'replace_control_characters\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'False\', \'None\'], "
+ }
+ member_method {
name: "unicode_encode"
argspec: "args=[\'input\', \'output_encoding\', \'errors\', \'replacement_char\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'None\'], "
}
@@ -61,6 +69,14 @@
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unicode_split"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'None\'], "
+ }
+ member_method {
+ name: "unicode_split_with_offsets"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'None\'], "
+ }
+ member_method {
name: "unicode_transcode"
argspec: "args=[\'input\', \'input_encoding\', \'output_encoding\', \'errors\', \'replacement_char\', \'replace_control_characters\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'False\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-attr-value.-list-value.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-attr-value.-list-value.pbtxt
deleted file mode 100644
index f1dffd5..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-attr-value.-list-value.pbtxt
+++ /dev/null
@@ -1,70 +0,0 @@
-path: "tensorflow.AttrValue.ListValue"
-tf_proto {
- descriptor {
- name: "ListValue"
- field {
- name: "s"
- number: 2
- label: LABEL_REPEATED
- type: TYPE_BYTES
- }
- field {
- name: "i"
- number: 3
- label: LABEL_REPEATED
- type: TYPE_INT64
- options {
- packed: true
- }
- }
- field {
- name: "f"
- number: 4
- label: LABEL_REPEATED
- type: TYPE_FLOAT
- options {
- packed: true
- }
- }
- field {
- name: "b"
- number: 5
- label: LABEL_REPEATED
- type: TYPE_BOOL
- options {
- packed: true
- }
- }
- field {
- name: "type"
- number: 6
- label: LABEL_REPEATED
- type: TYPE_ENUM
- type_name: ".tensorflow.DataType"
- options {
- packed: true
- }
- }
- field {
- name: "shape"
- number: 7
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.TensorShapeProto"
- }
- field {
- name: "tensor"
- number: 8
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.TensorProto"
- }
- field {
- name: "func"
- number: 9
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.NameAttrList"
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-attr-value.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-attr-value.pbtxt
deleted file mode 100644
index 6ccd64f..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-attr-value.pbtxt
+++ /dev/null
@@ -1,151 +0,0 @@
-path: "tensorflow.AttrValue"
-tf_proto {
- descriptor {
- name: "AttrValue"
- field {
- name: "s"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BYTES
- oneof_index: 0
- }
- field {
- name: "i"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- oneof_index: 0
- }
- field {
- name: "f"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_FLOAT
- oneof_index: 0
- }
- field {
- name: "b"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- oneof_index: 0
- }
- field {
- name: "type"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.DataType"
- oneof_index: 0
- }
- field {
- name: "shape"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.TensorShapeProto"
- oneof_index: 0
- }
- field {
- name: "tensor"
- number: 8
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.TensorProto"
- oneof_index: 0
- }
- field {
- name: "list"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.AttrValue.ListValue"
- oneof_index: 0
- }
- field {
- name: "func"
- number: 10
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.NameAttrList"
- oneof_index: 0
- }
- field {
- name: "placeholder"
- number: 9
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- oneof_index: 0
- }
- nested_type {
- name: "ListValue"
- field {
- name: "s"
- number: 2
- label: LABEL_REPEATED
- type: TYPE_BYTES
- }
- field {
- name: "i"
- number: 3
- label: LABEL_REPEATED
- type: TYPE_INT64
- options {
- packed: true
- }
- }
- field {
- name: "f"
- number: 4
- label: LABEL_REPEATED
- type: TYPE_FLOAT
- options {
- packed: true
- }
- }
- field {
- name: "b"
- number: 5
- label: LABEL_REPEATED
- type: TYPE_BOOL
- options {
- packed: true
- }
- }
- field {
- name: "type"
- number: 6
- label: LABEL_REPEATED
- type: TYPE_ENUM
- type_name: ".tensorflow.DataType"
- options {
- packed: true
- }
- }
- field {
- name: "shape"
- number: 7
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.TensorShapeProto"
- }
- field {
- name: "tensor"
- number: 8
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.TensorProto"
- }
- field {
- name: "func"
- number: 9
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.NameAttrList"
- }
- }
- oneof_decl {
- name: "value"
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.-device-count-entry.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.-device-count-entry.pbtxt
deleted file mode 100644
index d9b1426..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.-device-count-entry.pbtxt
+++ /dev/null
@@ -1,21 +0,0 @@
-path: "tensorflow.ConfigProto.DeviceCountEntry"
-tf_proto {
- descriptor {
- name: "DeviceCountEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- options {
- map_entry: true
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.-experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.-experimental.pbtxt
deleted file mode 100644
index caa72fe..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.-experimental.pbtxt
+++ /dev/null
@@ -1,34 +0,0 @@
-path: "tensorflow.ConfigProto.Experimental"
-tf_proto {
- descriptor {
- name: "Experimental"
- field {
- name: "collective_group_leader"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "executor_type"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "recv_buf_max_chunk"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "use_numa_affinity"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- reserved_range {
- start: 2
- end: 3
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.pbtxt
deleted file mode 100644
index b505d81..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-config-proto.pbtxt
+++ /dev/null
@@ -1,158 +0,0 @@
-path: "tensorflow.ConfigProto"
-tf_proto {
- descriptor {
- name: "ConfigProto"
- field {
- name: "device_count"
- number: 1
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.ConfigProto.DeviceCountEntry"
- }
- field {
- name: "intra_op_parallelism_threads"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "inter_op_parallelism_threads"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "use_per_session_threads"
- number: 9
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "session_inter_op_thread_pool"
- number: 12
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.ThreadPoolOptionProto"
- }
- field {
- name: "placement_period"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "device_filters"
- number: 4
- label: LABEL_REPEATED
- type: TYPE_STRING
- }
- field {
- name: "gpu_options"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.GPUOptions"
- }
- field {
- name: "allow_soft_placement"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "log_device_placement"
- number: 8
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "graph_options"
- number: 10
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.GraphOptions"
- }
- field {
- name: "operation_timeout_in_ms"
- number: 11
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "rpc_options"
- number: 13
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.RPCOptions"
- }
- field {
- name: "cluster_def"
- number: 14
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.ClusterDef"
- }
- field {
- name: "isolate_session_state"
- number: 15
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "experimental"
- number: 16
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.ConfigProto.Experimental"
- }
- nested_type {
- name: "DeviceCountEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- options {
- map_entry: true
- }
- }
- nested_type {
- name: "Experimental"
- field {
- name: "collective_group_leader"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "executor_type"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "recv_buf_max_chunk"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "use_numa_affinity"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- reserved_range {
- start: 2
- end: 3
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-g-p-u-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-g-p-u-options.pbtxt
deleted file mode 100644
index a2cc074..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-g-p-u-options.pbtxt
+++ /dev/null
@@ -1,98 +0,0 @@
-path: "tensorflow.GPUOptions"
-tf_proto {
- descriptor {
- name: "GPUOptions"
- field {
- name: "per_process_gpu_memory_fraction"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_DOUBLE
- }
- field {
- name: "allow_growth"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "allocator_type"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "deferred_deletion_bytes"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "visible_device_list"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "polling_active_delay_usecs"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "polling_inactive_delay_msecs"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "force_gpu_compatible"
- number: 8
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "experimental"
- number: 9
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.GPUOptions.Experimental"
- }
- nested_type {
- name: "Experimental"
- field {
- name: "virtual_devices"
- number: 1
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.GPUOptions.Experimental.VirtualDevices"
- }
- field {
- name: "use_unified_memory"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "num_dev_to_dev_copy_streams"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "collective_ring_order"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- nested_type {
- name: "VirtualDevices"
- field {
- name: "memory_limit_mb"
- number: 1
- label: LABEL_REPEATED
- type: TYPE_FLOAT
- }
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-graph-def.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-graph-def.pbtxt
deleted file mode 100644
index 19eccff..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-graph-def.pbtxt
+++ /dev/null
@@ -1,36 +0,0 @@
-path: "tensorflow.GraphDef"
-tf_proto {
- descriptor {
- name: "GraphDef"
- field {
- name: "node"
- number: 1
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.NodeDef"
- }
- field {
- name: "versions"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.VersionDef"
- }
- field {
- name: "version"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- options {
- deprecated: true
- }
- }
- field {
- name: "library"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.FunctionDefLibrary"
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-graph-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-graph-options.pbtxt
deleted file mode 100644
index a9f99bc..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-graph-options.pbtxt
+++ /dev/null
@@ -1,67 +0,0 @@
-path: "tensorflow.GraphOptions"
-tf_proto {
- descriptor {
- name: "GraphOptions"
- field {
- name: "enable_recv_scheduling"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "optimizer_options"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.OptimizerOptions"
- }
- field {
- name: "build_cost_model"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "build_cost_model_after"
- number: 9
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "infer_shapes"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "place_pruned_graph"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "enable_bfloat16_sendrecv"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "timeline_step"
- number: 8
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "rewrite_options"
- number: 10
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.RewriterConfig"
- }
- reserved_range {
- start: 1
- end: 2
- }
- reserved_name: "skip_common_subexpression_elimination"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-histogram-proto.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-histogram-proto.pbtxt
deleted file mode 100644
index d4402f3..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-histogram-proto.pbtxt
+++ /dev/null
@@ -1,54 +0,0 @@
-path: "tensorflow.HistogramProto"
-tf_proto {
- descriptor {
- name: "HistogramProto"
- field {
- name: "min"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_DOUBLE
- }
- field {
- name: "max"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_DOUBLE
- }
- field {
- name: "num"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_DOUBLE
- }
- field {
- name: "sum"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_DOUBLE
- }
- field {
- name: "sum_squares"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_DOUBLE
- }
- field {
- name: "bucket_limit"
- number: 6
- label: LABEL_REPEATED
- type: TYPE_DOUBLE
- options {
- packed: true
- }
- }
- field {
- name: "bucket"
- number: 7
- label: LABEL_REPEATED
- type: TYPE_DOUBLE
- options {
- packed: true
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-log-message.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-log-message.pbtxt
deleted file mode 100644
index 5023aa9..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-log-message.pbtxt
+++ /dev/null
@@ -1,46 +0,0 @@
-path: "tensorflow.LogMessage"
-tf_proto {
- descriptor {
- name: "LogMessage"
- field {
- name: "level"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.LogMessage.Level"
- }
- field {
- name: "message"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- enum_type {
- name: "Level"
- value {
- name: "UNKNOWN"
- number: 0
- }
- value {
- name: "DEBUGGING"
- number: 10
- }
- value {
- name: "INFO"
- number: 20
- }
- value {
- name: "WARN"
- number: 30
- }
- value {
- name: "ERROR"
- number: 40
- }
- value {
- name: "FATAL"
- number: 50
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
deleted file mode 100644
index 0ba09be..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
+++ /dev/null
@@ -1,22 +0,0 @@
-path: "tensorflow.MetaGraphDef.CollectionDefEntry"
-tf_proto {
- descriptor {
- name: "CollectionDefEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.CollectionDef"
- }
- options {
- map_entry: true
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-meta-info-def.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
deleted file mode 100644
index 41c62a4..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
+++ /dev/null
@@ -1,50 +0,0 @@
-path: "tensorflow.MetaGraphDef.MetaInfoDef"
-tf_proto {
- descriptor {
- name: "MetaInfoDef"
- field {
- name: "meta_graph_version"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "stripped_op_list"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.OpList"
- }
- field {
- name: "any_info"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".google.protobuf.Any"
- }
- field {
- name: "tags"
- number: 4
- label: LABEL_REPEATED
- type: TYPE_STRING
- }
- field {
- name: "tensorflow_version"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "tensorflow_git_version"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "stripped_default_attrs"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
deleted file mode 100644
index 73dc414..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
+++ /dev/null
@@ -1,22 +0,0 @@
-path: "tensorflow.MetaGraphDef.SignatureDefEntry"
-tf_proto {
- descriptor {
- name: "SignatureDefEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.SignatureDef"
- }
- options {
- map_entry: true
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.pbtxt
deleted file mode 100644
index d71c235..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-meta-graph-def.pbtxt
+++ /dev/null
@@ -1,133 +0,0 @@
-path: "tensorflow.MetaGraphDef"
-tf_proto {
- descriptor {
- name: "MetaGraphDef"
- field {
- name: "meta_info_def"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.MetaGraphDef.MetaInfoDef"
- }
- field {
- name: "graph_def"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.GraphDef"
- }
- field {
- name: "saver_def"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.SaverDef"
- }
- field {
- name: "collection_def"
- number: 4
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.MetaGraphDef.CollectionDefEntry"
- }
- field {
- name: "signature_def"
- number: 5
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.MetaGraphDef.SignatureDefEntry"
- }
- field {
- name: "asset_file_def"
- number: 6
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.AssetFileDef"
- }
- nested_type {
- name: "MetaInfoDef"
- field {
- name: "meta_graph_version"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "stripped_op_list"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.OpList"
- }
- field {
- name: "any_info"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".google.protobuf.Any"
- }
- field {
- name: "tags"
- number: 4
- label: LABEL_REPEATED
- type: TYPE_STRING
- }
- field {
- name: "tensorflow_version"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "tensorflow_git_version"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "stripped_default_attrs"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- }
- nested_type {
- name: "CollectionDefEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.CollectionDef"
- }
- options {
- map_entry: true
- }
- }
- nested_type {
- name: "SignatureDefEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.SignatureDef"
- }
- options {
- map_entry: true
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-name-attr-list.-attr-entry.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-name-attr-list.-attr-entry.pbtxt
deleted file mode 100644
index b119b20..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-name-attr-list.-attr-entry.pbtxt
+++ /dev/null
@@ -1,22 +0,0 @@
-path: "tensorflow.NameAttrList.AttrEntry"
-tf_proto {
- descriptor {
- name: "AttrEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.AttrValue"
- }
- options {
- map_entry: true
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-name-attr-list.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-name-attr-list.pbtxt
deleted file mode 100644
index fcdb411..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-name-attr-list.pbtxt
+++ /dev/null
@@ -1,38 +0,0 @@
-path: "tensorflow.NameAttrList"
-tf_proto {
- descriptor {
- name: "NameAttrList"
- field {
- name: "name"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "attr"
- number: 2
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.NameAttrList.AttrEntry"
- }
- nested_type {
- name: "AttrEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.AttrValue"
- }
- options {
- map_entry: true
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-node-def.-attr-entry.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-node-def.-attr-entry.pbtxt
deleted file mode 100644
index 622e4c3..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-node-def.-attr-entry.pbtxt
+++ /dev/null
@@ -1,22 +0,0 @@
-path: "tensorflow.NodeDef.AttrEntry"
-tf_proto {
- descriptor {
- name: "AttrEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.AttrValue"
- }
- options {
- map_entry: true
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-node-def.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-node-def.pbtxt
deleted file mode 100644
index 646fa8a..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-node-def.pbtxt
+++ /dev/null
@@ -1,56 +0,0 @@
-path: "tensorflow.NodeDef"
-tf_proto {
- descriptor {
- name: "NodeDef"
- field {
- name: "name"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "op"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "input"
- number: 3
- label: LABEL_REPEATED
- type: TYPE_STRING
- }
- field {
- name: "device"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "attr"
- number: 5
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.NodeDef.AttrEntry"
- }
- nested_type {
- name: "AttrEntry"
- field {
- name: "key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "value"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.AttrValue"
- }
- options {
- map_entry: true
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-optimizer-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-optimizer-options.pbtxt
deleted file mode 100644
index 3ccf9d4..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-optimizer-options.pbtxt
+++ /dev/null
@@ -1,74 +0,0 @@
-path: "tensorflow.OptimizerOptions"
-tf_proto {
- descriptor {
- name: "OptimizerOptions"
- field {
- name: "do_common_subexpression_elimination"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "do_constant_folding"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "max_folded_constant_in_bytes"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "do_function_inlining"
- number: 4
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "opt_level"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.OptimizerOptions.Level"
- }
- field {
- name: "global_jit_level"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.OptimizerOptions.GlobalJitLevel"
- }
- enum_type {
- name: "Level"
- value {
- name: "L1"
- number: 0
- }
- value {
- name: "L0"
- number: -1
- }
- }
- enum_type {
- name: "GlobalJitLevel"
- value {
- name: "DEFAULT"
- number: 0
- }
- value {
- name: "OFF"
- number: -1
- }
- value {
- name: "ON_1"
- number: 1
- }
- value {
- name: "ON_2"
- number: 2
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-ragged-tensor.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-ragged-tensor.pbtxt
new file mode 100644
index 0000000..c0ed956
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.-ragged-tensor.pbtxt
@@ -0,0 +1,125 @@
+path: "tensorflow.RaggedTensor"
+tf_class {
+ is_instance: "<class \'tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "dtype"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "flat_values"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "nested_row_splits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "ragged_rank"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "row_splits"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "shape"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "values"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'values\', \'row_splits\', \'cached_row_lengths\', \'cached_value_rowids\', \'cached_nrows\', \'internal\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'False\'], "
+ }
+ member_method {
+ name: "bounding_shape"
+ argspec: "args=[\'self\', \'axis\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_nested_row_lengths"
+ argspec: "args=[\'cls\', \'flat_values\', \'nested_row_lengths\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_nested_row_splits"
+ argspec: "args=[\'cls\', \'flat_values\', \'nested_row_splits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_nested_value_rowids"
+ argspec: "args=[\'cls\', \'flat_values\', \'nested_value_rowids\', \'nested_nrows\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "from_row_lengths"
+ argspec: "args=[\'cls\', \'values\', \'row_lengths\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_row_limits"
+ argspec: "args=[\'cls\', \'values\', \'row_limits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_row_splits"
+ argspec: "args=[\'cls\', \'values\', \'row_splits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_row_starts"
+ argspec: "args=[\'cls\', \'values\', \'row_starts\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_sparse"
+ argspec: "args=[\'cls\', \'st_input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "from_tensor"
+ argspec: "args=[\'cls\', \'tensor\', \'lengths\', \'padding\', \'ragged_rank\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'1\', \'None\'], "
+ }
+ member_method {
+ name: "from_value_rowids"
+ argspec: "args=[\'cls\', \'values\', \'value_rowids\', \'nrows\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "nested_row_lengths"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "nrows"
+ argspec: "args=[\'self\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'int64\'>\", \'None\'], "
+ }
+ member_method {
+ name: "row_lengths"
+ argspec: "args=[\'self\', \'axis\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "row_limits"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "row_starts"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "to_list"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "to_sparse"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "to_tensor"
+ argspec: "args=[\'self\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "value_rowids"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "with_flat_values"
+ argspec: "args=[\'self\', \'new_values\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "with_values"
+ argspec: "args=[\'self\', \'new_values\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-run-metadata.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-run-metadata.pbtxt
deleted file mode 100644
index 1287940..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-run-metadata.pbtxt
+++ /dev/null
@@ -1,27 +0,0 @@
-path: "tensorflow.RunMetadata"
-tf_proto {
- descriptor {
- name: "RunMetadata"
- field {
- name: "step_stats"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.StepStats"
- }
- field {
- name: "cost_graph"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.CostGraphDef"
- }
- field {
- name: "partition_graphs"
- number: 3
- label: LABEL_REPEATED
- type: TYPE_MESSAGE
- type_name: ".tensorflow.GraphDef"
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-run-options.-experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-run-options.-experimental.pbtxt
deleted file mode 100644
index 47b5b56..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-run-options.-experimental.pbtxt
+++ /dev/null
@@ -1,18 +0,0 @@
-path: "tensorflow.RunOptions.Experimental"
-tf_proto {
- descriptor {
- name: "Experimental"
- field {
- name: "collective_graph_key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "use_run_handler_pool"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-run-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-run-options.pbtxt
deleted file mode 100644
index c0c2e7b..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-run-options.pbtxt
+++ /dev/null
@@ -1,89 +0,0 @@
-path: "tensorflow.RunOptions"
-tf_proto {
- descriptor {
- name: "RunOptions"
- field {
- name: "trace_level"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.RunOptions.TraceLevel"
- }
- field {
- name: "timeout_in_ms"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "inter_op_thread_pool"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_INT32
- }
- field {
- name: "output_partition_graphs"
- number: 5
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "debug_options"
- number: 6
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.DebugOptions"
- }
- field {
- name: "report_tensor_allocations_upon_oom"
- number: 7
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- field {
- name: "experimental"
- number: 8
- label: LABEL_OPTIONAL
- type: TYPE_MESSAGE
- type_name: ".tensorflow.RunOptions.Experimental"
- }
- nested_type {
- name: "Experimental"
- field {
- name: "collective_graph_key"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_INT64
- }
- field {
- name: "use_run_handler_pool"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_BOOL
- }
- }
- enum_type {
- name: "TraceLevel"
- value {
- name: "NO_TRACE"
- number: 0
- }
- value {
- name: "SOFTWARE_TRACE"
- number: 1
- }
- value {
- name: "HARDWARE_TRACE"
- number: 2
- }
- value {
- name: "FULL_TRACE"
- number: 3
- }
- }
- reserved_range {
- start: 4
- end: 5
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-session-log.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-session-log.pbtxt
deleted file mode 100644
index 259f241..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-session-log.pbtxt
+++ /dev/null
@@ -1,44 +0,0 @@
-path: "tensorflow.SessionLog"
-tf_proto {
- descriptor {
- name: "SessionLog"
- field {
- name: "status"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.SessionLog.SessionStatus"
- }
- field {
- name: "checkpoint_path"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "msg"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- enum_type {
- name: "SessionStatus"
- value {
- name: "STATUS_UNSPECIFIED"
- number: 0
- }
- value {
- name: "START"
- number: 1
- }
- value {
- name: "STOP"
- number: 2
- }
- value {
- name: "CHECKPOINT"
- number: 3
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-sparse-tensor-value.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-sparse-tensor-value.pbtxt
deleted file mode 100644
index d33fd4d..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.-sparse-tensor-value.pbtxt
+++ /dev/null
@@ -1,26 +0,0 @@
-path: "tensorflow.SparseTensorValue"
-tf_class {
- is_instance: "<class \'tensorflow.python.framework.sparse_tensor.SparseTensorValue\'>"
- is_instance: "<type \'tuple\'>"
- member {
- name: "dense_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "indices"
- mtype: "<type \'property\'>"
- }
- member {
- name: "values"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- }
- member_method {
- name: "count"
- }
- member_method {
- name: "index"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt
index e85949f..6136c8f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.-variable.pbtxt
@@ -64,6 +64,10 @@
argspec: "args=[\'self\', \'delta\', \'use_locking\', \'name\', \'read_value\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'True\'], "
}
member_method {
+ name: "batch_scatter_update"
+ argspec: "args=[\'self\', \'sparse_delta\', \'use_locking\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ }
+ member_method {
name: "count_up_to"
argspec: "args=[\'self\', \'limit\'], varargs=None, keywords=None, defaults=None"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.-feature.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.-feature.pbtxt
new file mode 100644
index 0000000..a71da11
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.-feature.pbtxt
@@ -0,0 +1,28 @@
+path: "tensorflow.autograph.experimental.Feature"
+tf_class {
+ is_instance: "<enum \'Feature\'>"
+ member {
+ name: "ALL"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "AUTO_CONTROL_DEPS"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "DECORATORS"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "ERROR_REWRITING"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "LISTS"
+ mtype: "<enum \'Feature\'>"
+ }
+ member {
+ name: "NAME_SCOPES"
+ mtype: "<enum \'Feature\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.-verbosity.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.-verbosity.pbtxt
new file mode 100644
index 0000000..c4d5b77
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.-verbosity.pbtxt
@@ -0,0 +1,12 @@
+path: "tensorflow.autograph.experimental.Verbosity"
+tf_class {
+ is_instance: "<enum \'Verbosity\'>"
+ member {
+ name: "BRIEF"
+ mtype: "<enum \'Verbosity\'>"
+ }
+ member {
+ name: "VERBOSE"
+ mtype: "<enum \'Verbosity\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.pbtxt
new file mode 100644
index 0000000..5747dac
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.autograph.experimental.pbtxt
@@ -0,0 +1,11 @@
+path: "tensorflow.autograph.experimental"
+tf_module {
+ member {
+ name: "Feature"
+ mtype: "<class \'enum.EnumMeta\'>"
+ }
+ member {
+ name: "Verbosity"
+ mtype: "<class \'enum.EnumMeta\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.autograph.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.autograph.pbtxt
new file mode 100644
index 0000000..34bdab9
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.autograph.pbtxt
@@ -0,0 +1,15 @@
+path: "tensorflow.autograph"
+tf_module {
+ member {
+ name: "experimental"
+ mtype: "<type \'module\'>"
+ }
+ member_method {
+ name: "to_code"
+ argspec: "args=[\'entity\', \'recursive\', \'arg_values\', \'arg_types\', \'indentation\', \'optional_features\', \'experimental_partial_types\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\', \' \', \'Feature.ALL\', \'None\'], "
+ }
+ member_method {
+ name: "to_graph"
+ argspec: "args=[\'entity\', \'recursive\', \'verbose\', \'arg_values\', \'arg_types\', \'strip_decorators\', \'optional_features\', \'experimental_partial_types\'], varargs=None, keywords=None, defaults=[\'True\', \'Verbosity.VERBOSE\', \'None\', \'None\', \'None\', \'Feature.ALL\', \'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
index ac8dd2d..d877339 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.-dataset.pbtxt
@@ -4,15 +4,15 @@
is_instance: "<type \'object\'>"
member {
name: "output_classes"
- mtype: "<class \'abc.abstractproperty\'>"
+ mtype: "<type \'property\'>"
}
member {
name: "output_shapes"
- mtype: "<class \'abc.abstractproperty\'>"
+ mtype: "<type \'property\'>"
}
member {
name: "output_types"
- mtype: "<class \'abc.abstractproperty\'>"
+ mtype: "<type \'property\'>"
}
member_method {
name: "__init__"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-dataset-structure.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-dataset-structure.pbtxt
new file mode 100644
index 0000000..dcb304f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-dataset-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.DatasetStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.ops.dataset_ops.DatasetStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'element_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-nested-structure.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-nested-structure.pbtxt
new file mode 100644
index 0000000..b4b066e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-nested-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.NestedStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.NestedStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'nested_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt
index 9ca7582..3b7ad64 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optimization-options.pbtxt
@@ -4,6 +4,10 @@
is_instance: "<class \'tensorflow.python.data.util.options.OptionsBase\'>"
is_instance: "<type \'object\'>"
member {
+ name: "apply_default_optimizations"
+ mtype: "<type \'property\'>"
+ }
+ member {
name: "filter_fusion"
mtype: "<type \'property\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optional-structure.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optional-structure.pbtxt
new file mode 100644
index 0000000..bf41c1d
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-optional-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.OptionalStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.ops.optional_ops.OptionalStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'value_structure\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sparse-tensor-structure.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sparse-tensor-structure.pbtxt
new file mode 100644
index 0000000..f97376b
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-sparse-tensor-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.SparseTensorStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.SparseTensorStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtype\', \'dense_shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-structure.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-structure.pbtxt
new file mode 100644
index 0000000..a99db45
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-structure.pbtxt
@@ -0,0 +1,16 @@
+path: "tensorflow.data.experimental.Structure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-tensor-structure.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-tensor-structure.pbtxt
new file mode 100644
index 0000000..f5c8864
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-tensor-structure.pbtxt
@@ -0,0 +1,18 @@
+path: "tensorflow.data.experimental.TensorStructure"
+tf_class {
+ is_instance: "<class \'tensorflow.python.data.util.structure.TensorStructure\'>"
+ is_instance: "<class \'tensorflow.python.data.util.structure.Structure\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'dtype\', \'shape\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "from_value"
+ argspec: "args=[\'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "is_compatible_with"
+ argspec: "args=[\'self\', \'other\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt
index f981b1a..2d11590 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.pbtxt
@@ -13,10 +13,18 @@
mtype: "<type \'type\'>"
}
member {
+ name: "DatasetStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "INFINITE_CARDINALITY"
mtype: "<type \'int\'>"
}
member {
+ name: "NestedStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "OptimizationOptions"
mtype: "<type \'type\'>"
}
@@ -25,6 +33,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "OptionalStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "RandomDataset"
mtype: "<type \'type\'>"
}
@@ -33,6 +45,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "SparseTensorStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "SqlDataset"
mtype: "<type \'type\'>"
}
@@ -45,10 +61,18 @@
mtype: "<type \'type\'>"
}
member {
+ name: "Structure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "TFRecordWriter"
mtype: "<type \'type\'>"
}
member {
+ name: "TensorStructure"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "ThreadingOptions"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
new file mode 100644
index 0000000..81224f0
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
@@ -0,0 +1,142 @@
+path: "tensorflow.distribute.MirroredStrategy"
+tf_class {
+ is_instance: "<class \'tensorflow.python.distribute.mirrored_strategy.MirroredStrategy\'>"
+ is_instance: "<class \'tensorflow.python.distribute.distribute_lib.DistributionStrategy\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "between_graph"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "extended"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "num_replicas_in_sync"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "parameter_devices"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "require_static_shapes"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "should_checkpoint"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "should_init"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "should_save_summary"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "worker_devices"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'devices\', \'cross_device_ops\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "batch_reduce"
+ argspec: "args=[\'self\', \'aggregation\', \'value_destination_pairs\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "broadcast"
+ argspec: "args=[\'self\', \'tensor\', \'destinations\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "call_for_each_replica"
+ argspec: "args=[\'self\', \'fn\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "colocate_vars_with"
+ argspec: "args=[\'self\', \'colocate_with_variable\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "configure"
+ argspec: "args=[\'self\', \'session_config\', \'cluster_spec\', \'task_type\', \'task_id\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "distribute_dataset"
+ argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "experimental_finalize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "experimental_initialize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "experimental_run"
+ argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "finalize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "group"
+ argspec: "args=[\'self\', \'value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "initialize"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "make_dataset_iterator"
+ argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "make_input_fn_iterator"
+ argspec: "args=[\'self\', \'input_fn\', \'replication_mode\'], varargs=None, keywords=None, defaults=[\'InputReplicationMode.PER_WORKER\'], "
+ }
+ member_method {
+ name: "non_slot_devices"
+ argspec: "args=[\'self\', \'var_list\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "read_var"
+ argspec: "args=[\'self\', \'v\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reduce"
+ argspec: "args=[\'self\', \'reduce_op\', \'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "run_steps_on_dataset"
+ argspec: "args=[\'self\', \'fn\', \'iterator\', \'iterations\', \'initial_loop_values\'], varargs=None, keywords=None, defaults=[\'1\', \'None\'], "
+ }
+ member_method {
+ name: "scope"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "unwrap"
+ argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "update"
+ argspec: "args=[\'self\', \'var\', \'fn\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "update_config_proto"
+ argspec: "args=[\'self\', \'config_proto\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "update_non_slot"
+ argspec: "args=[\'self\', \'colocate_with\', \'fn\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "value_container"
+ argspec: "args=[\'self\', \'value\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-server.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-server.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-server.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.distribute.-server.pbtxt
index 9b8f185..6c39bf4 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-server.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-server.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.Server"
+path: "tensorflow.distribute.Server"
tf_class {
is_instance: "<class \'tensorflow.python.training.server_lib.Server\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
index 9eb73d2..63b6584 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
@@ -75,6 +75,10 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "experimental_run"
+ argspec: "args=[\'self\', \'fn\', \'input_iterator\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "finalize"
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 4d833b5..31dc6e0 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
@@ -9,6 +9,10 @@
mtype: "<class \'enum.EnumMeta\'>"
}
member {
+ name: "MirroredStrategy"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "ReduceOp"
mtype: "<class \'enum.EnumMeta\'>"
}
@@ -17,6 +21,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "Server"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "Strategy"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt
index 848fc30..01b870a8 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.dtypes.pbtxt
@@ -45,6 +45,10 @@
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
member {
+ name: "int16"
+ mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
+ }
+ member {
name: "int32"
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-checkpoint-saver-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-checkpoint-saver-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-checkpoint-saver-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-checkpoint-saver-hook.pbtxt
index c3037ba..f9e1504 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-checkpoint-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-checkpoint-saver-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.CheckpointSaverHook"
+path: "tensorflow.estimator.CheckpointSaverHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.CheckpointSaverHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.-checkpoint-saver-listener.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-checkpoint-saver-listener.pbtxt
new file mode 100644
index 0000000..111b758
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-checkpoint-saver-listener.pbtxt
@@ -0,0 +1,24 @@
+path: "tensorflow.estimator.CheckpointSaverListener"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.CheckpointSaverListener\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "after_save"
+ argspec: "args=[\'self\', \'session\', \'global_step_value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "before_save"
+ argspec: "args=[\'self\', \'session\', \'global_step_value\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "begin"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "end"
+ argspec: "args=[\'self\', \'session\', \'global_step_value\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-feed-fn-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-feed-fn-hook.pbtxt
similarity index 96%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-feed-fn-hook.pbtxt
rename to tensorflow/tools/api/golden/v2/tensorflow.estimator.-feed-fn-hook.pbtxt
index 7bec4d0..f24de49 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-feed-fn-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-feed-fn-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.FeedFnHook"
+path: "tensorflow.estimator.FeedFnHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.FeedFnHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-final-ops-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-final-ops-hook.pbtxt
similarity index 96%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-final-ops-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-final-ops-hook.pbtxt
index 31cf9aa..6651170 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-final-ops-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-final-ops-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.FinalOpsHook"
+path: "tensorflow.estimator.FinalOpsHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.FinalOpsHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-global-step-waiter-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-global-step-waiter-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-global-step-waiter-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-global-step-waiter-hook.pbtxt
index 1474486..37db48b 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-global-step-waiter-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-global-step-waiter-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.GlobalStepWaiterHook"
+path: "tensorflow.estimator.GlobalStepWaiterHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.GlobalStepWaiterHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-logging-tensor-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-logging-tensor-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-logging-tensor-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-logging-tensor-hook.pbtxt
index 9801c05..425f016 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-logging-tensor-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-logging-tensor-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.LoggingTensorHook"
+path: "tensorflow.estimator.LoggingTensorHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.LoggingTensorHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.-nan-loss-during-training-error.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-nan-loss-during-training-error.pbtxt
new file mode 100644
index 0000000..6cf6e17
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-nan-loss-during-training-error.pbtxt
@@ -0,0 +1,12 @@
+path: "tensorflow.estimator.NanLossDuringTrainingError"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.NanLossDuringTrainingError\'>"
+ is_instance: "<type \'exceptions.RuntimeError\'>"
+ member {
+ name: "args"
+ mtype: "<type \'getset_descriptor\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-nan-tensor-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-nan-tensor-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-nan-tensor-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-nan-tensor-hook.pbtxt
index 7d1c89f..82293c2 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-nan-tensor-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-nan-tensor-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.NanTensorHook"
+path: "tensorflow.estimator.NanTensorHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.NanTensorHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-profiler-hook.pbtxt
similarity index 77%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-profiler-hook.pbtxt
index e388599..65b5fb1 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-profiler-hook.pbtxt
@@ -1,11 +1,11 @@
-path: "tensorflow.train.StopAtStepHook"
+path: "tensorflow.estimator.ProfilerHook"
tf_class {
- is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.StopAtStepHook\'>"
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.ProfilerHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
is_instance: "<type \'object\'>"
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'num_steps\', \'last_step\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ argspec: "args=[\'self\', \'save_steps\', \'save_secs\', \'output_dir\', \'show_dataflow\', \'show_memory\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'\', \'True\', \'False\'], "
}
member_method {
name: "after_create_session"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.-second-or-step-timer.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-second-or-step-timer.pbtxt
new file mode 100644
index 0000000..64051d2
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-second-or-step-timer.pbtxt
@@ -0,0 +1,26 @@
+path: "tensorflow.estimator.SecondOrStepTimer"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.SecondOrStepTimer\'>"
+ is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks._HookTimer\'>"
+ is_instance: "<type \'object\'>"
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'every_secs\', \'every_steps\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
+ name: "last_triggered_step"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "reset"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "should_trigger_for_step"
+ argspec: "args=[\'self\', \'step\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "update_last_triggered_step"
+ argspec: "args=[\'self\', \'step\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-args.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-args.pbtxt
new file mode 100644
index 0000000..b375c74
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-args.pbtxt
@@ -0,0 +1,27 @@
+path: "tensorflow.estimator.SessionRunArgs"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunArgs\'>"
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunArgs\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "feed_dict"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "fetches"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "options"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-context.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-context.pbtxt
new file mode 100644
index 0000000..cb4ac9f
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-context.pbtxt
@@ -0,0 +1,25 @@
+path: "tensorflow.estimator.SessionRunContext"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunContext\'>"
+ is_instance: "<type \'object\'>"
+ member {
+ name: "original_args"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "session"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "stop_requested"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ argspec: "args=[\'self\', \'original_args\', \'session\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "request_stop"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-session-run-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-hook.pbtxt
similarity index 94%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-session-run-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-hook.pbtxt
index db1aa24..54e9ad9 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-session-run-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.SessionRunHook"
+path: "tensorflow.estimator.SessionRunHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-values.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-values.pbtxt
new file mode 100644
index 0000000..6788141
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-session-run-values.pbtxt
@@ -0,0 +1,27 @@
+path: "tensorflow.estimator.SessionRunValues"
+tf_class {
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunValues\'>"
+ is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunValues\'>"
+ is_instance: "<type \'tuple\'>"
+ member {
+ name: "options"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "results"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "run_metadata"
+ mtype: "<type \'property\'>"
+ }
+ member_method {
+ name: "__init__"
+ }
+ member_method {
+ name: "count"
+ }
+ member_method {
+ name: "index"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-step-counter-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-step-counter-hook.pbtxt
similarity index 96%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-step-counter-hook.pbtxt
rename to tensorflow/tools/api/golden/v2/tensorflow.estimator.-step-counter-hook.pbtxt
index 13261f6..4368e04 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-step-counter-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-step-counter-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.StepCounterHook"
+path: "tensorflow.estimator.StepCounterHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.StepCounterHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-stop-at-step-hook.pbtxt
similarity index 95%
copy from tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.estimator.-stop-at-step-hook.pbtxt
index e388599..938b189 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-stop-at-step-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-stop-at-step-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.StopAtStepHook"
+path: "tensorflow.estimator.StopAtStepHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.StopAtStepHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-summary-saver-hook.pbtxt
similarity index 96%
rename from tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt
rename to tensorflow/tools/api/golden/v2/tensorflow.estimator.-summary-saver-hook.pbtxt
index 697c366..1041573 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.-summary-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.-summary-saver-hook.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.train.SummarySaverHook"
+path: "tensorflow.estimator.SummarySaverHook"
tf_class {
is_instance: "<class \'tensorflow.python.training.basic_session_run_hooks.SummarySaverHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt
index aba1202..5a2a01c 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.experimental.-in-memory-evaluator-hook.pbtxt
@@ -1,6 +1,6 @@
path: "tensorflow.estimator.experimental.InMemoryEvaluatorHook"
tf_class {
- is_instance: "<class \'tensorflow_estimator.python.estimator.hooks.InMemoryEvaluatorHook\'>"
+ is_instance: "<class \'tensorflow_estimator.python.estimator.hooks.hooks.InMemoryEvaluatorHook\'>"
is_instance: "<class \'tensorflow.python.training.session_run_hook.SessionRunHook\'>"
is_instance: "<type \'object\'>"
member_method {
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.estimator.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.estimator.pbtxt
index c5b0085..6f57505 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.estimator.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.estimator.pbtxt
@@ -25,6 +25,14 @@
mtype: "<type \'type\'>"
}
member {
+ name: "CheckpointSaverHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "CheckpointSaverListener"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "DNNClassifier"
mtype: "<type \'type\'>"
}
@@ -65,10 +73,22 @@
mtype: "<type \'type\'>"
}
member {
+ name: "FeedFnHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "FinalExporter"
mtype: "<type \'type\'>"
}
member {
+ name: "FinalOpsHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "GlobalStepWaiterHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "LatestExporter"
mtype: "<type \'type\'>"
}
@@ -85,14 +105,62 @@
mtype: "<type \'type\'>"
}
member {
+ name: "LoggingTensorHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "ModeKeys"
mtype: "<type \'type\'>"
}
member {
+ name: "NanLossDuringTrainingError"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "NanTensorHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "ProfilerHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "RunConfig"
mtype: "<type \'type\'>"
}
member {
+ name: "SecondOrStepTimer"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunArgs"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunContext"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SessionRunValues"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "StepCounterHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "StopAtStepHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "SummarySaverHook"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "TrainSpec"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.gfile.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.gfile.pbtxt
index 93d9b0f..cfa3372 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.gfile.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.gfile.pbtxt
@@ -46,6 +46,6 @@
}
member_method {
name: "walk"
- argspec: "args=[\'top\', \'topdown\', \'onerror\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ argspec: "args=[\'top\', \'topdown\', \'onerror\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
}
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt
index 8906329..2d9c759 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.io.pbtxt
@@ -9,22 +9,6 @@
mtype: "<type \'type\'>"
}
member {
- name: "PaddingFIFOQueue"
- mtype: "<type \'type\'>"
- }
- member {
- name: "PriorityQueue"
- mtype: "<type \'type\'>"
- }
- member {
- name: "QueueBase"
- mtype: "<type \'type\'>"
- }
- member {
- name: "RandomShuffleQueue"
- mtype: "<type \'type\'>"
- }
- member {
name: "SparseFeature"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-cu-d-n-n-l-s-t-m.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-cu-d-n-n-l-s-t-m.pbtxt
deleted file mode 100644
index 7c463ff..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-cu-d-n-n-l-s-t-m.pbtxt
+++ /dev/null
@@ -1,197 +0,0 @@
-path: "tensorflow.keras.layers.CuDNNLSTM"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.layers.cudnn_recurrent.CuDNNLSTM\'>"
- is_instance: "<class \'tensorflow.python.keras.layers.cudnn_recurrent._CuDNNRNN\'>"
- is_instance: "<class \'tensorflow.python.keras.layers.recurrent.RNN\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "cell"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "states"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'units\', \'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\'], varargs=None, keywords=kwargs, defaults=[\'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'False\', \'False\', \'False\', \'False\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'synchronization\', \'aggregation\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\', \'mask\', \'training\', \'initial_state\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_initial_state"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\', \'states\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-l-s-t-m.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-l-s-t-m.pbtxt
index 529c750..9144a5b 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-l-s-t-m.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-l-s-t-m.pbtxt
@@ -1,5 +1,6 @@
path: "tensorflow.keras.layers.LSTM"
tf_class {
+ is_instance: "<class \'tensorflow.python.keras.layers.recurrent.UnifiedLSTM\'>"
is_instance: "<class \'tensorflow.python.keras.layers.recurrent.LSTM\'>"
is_instance: "<class \'tensorflow.python.keras.layers.recurrent.RNN\'>"
is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
@@ -155,7 +156,7 @@
}
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'units\', \'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\', \'dropout\', \'recurrent_dropout\', \'implementation\', \'return_sequences\', \'return_state\', \'go_backwards\', \'stateful\', \'unroll\'], varargs=None, keywords=kwargs, defaults=[\'tanh\', \'hard_sigmoid\', \'True\', \'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'0.0\', \'0.0\', \'1\', \'False\', \'False\', \'False\', \'False\', \'False\'], "
+ argspec: "args=[\'self\', \'units\', \'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\', \'dropout\', \'recurrent_dropout\', \'implementation\', \'return_sequences\', \'return_state\', \'go_backwards\', \'stateful\', \'time_major\', \'unroll\'], varargs=None, keywords=kwargs, defaults=[\'tanh\', \'sigmoid\', \'True\', \'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'0.0\', \'0.0\', \'1\', \'False\', \'False\', \'False\', \'False\', \'False\', \'False\'], "
}
member_method {
name: "add_loss"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.pbtxt
index 3b4724e..10ac3a7 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.pbtxt
@@ -117,10 +117,6 @@
mtype: "<type \'type\'>"
}
member {
- name: "CuDNNLSTM"
- mtype: "<type \'type\'>"
- }
- member {
name: "Dense"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.-sensitivity-at-specificity.pbtxt
similarity index 92%
copy from tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.-sensitivity-at-specificity.pbtxt
index 33226a2..aa77d19 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.-sensitivity-at-specificity.pbtxt
@@ -1,7 +1,7 @@
-path: "tensorflow.metrics.FalseNegatives"
+path: "tensorflow.keras.metrics.SensitivityAtSpecificity"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalseNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SensitivityAtSpecificity\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SensitivitySpecificityBase\'>"
is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
@@ -84,7 +84,7 @@
}
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ argspec: "args=[\'self\', \'specificity\', \'num_thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'200\', \'None\', \'None\'], "
}
member_method {
name: "add_loss"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.-specificity-at-sensitivity.pbtxt
similarity index 92%
copy from tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.-specificity-at-sensitivity.pbtxt
index 33226a2..67857aa 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.metrics.-false-negatives.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.-specificity-at-sensitivity.pbtxt
@@ -1,7 +1,7 @@
-path: "tensorflow.metrics.FalseNegatives"
+path: "tensorflow.keras.metrics.SpecificityAtSensitivity"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalseNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SpecificityAtSensitivity\'>"
+ is_instance: "<class \'tensorflow.python.keras.metrics.SensitivitySpecificityBase\'>"
is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
@@ -84,7 +84,7 @@
}
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
+ argspec: "args=[\'self\', \'sensitivity\', \'num_thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'200\', \'None\', \'None\'], "
}
member_method {
name: "add_loss"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.pbtxt
index 8cab17e..905021d 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.metrics.pbtxt
@@ -33,10 +33,18 @@
mtype: "<type \'type\'>"
}
member {
+ name: "SensitivityAtSpecificity"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "SparseCategoricalAccuracy"
mtype: "<type \'type\'>"
}
member {
+ name: "SpecificityAtSensitivity"
+ mtype: "<type \'type\'>"
+ }
+ member {
name: "TrueNegatives"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adadelta.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adadelta.pbtxt
index b9ce154..5426269 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adadelta.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adadelta.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.Adadelta"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adadelta\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adadelta.Adadelta\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'rho\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'1.0\', \'0.95\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'rho\', \'epsilon\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.95\', \'1e-07\', \'Adadelta\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adagrad.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adagrad.pbtxt
index d0dc9e3..c39fe6b 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adagrad.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adagrad.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.Adagrad"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adagrad\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adagrad.Adagrad\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'0.01\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'initial_accumulator_value\', \'epsilon\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.1\', \'1e-07\', \'Adagrad\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adam.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adam.pbtxt
index 06815fa..05d46d3 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adam.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adam.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.Adam"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adam\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adam.Adam\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'beta_1\', \'beta_2\', \'epsilon\', \'decay\', \'amsgrad\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.999\', \'None\', \'0.0\', \'False\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'beta_1\', \'beta_2\', \'epsilon\', \'amsgrad\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.999\', \'1e-07\', \'False\', \'Adam\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adamax.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adamax.pbtxt
index 47b55fd..78829de 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adamax.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-adamax.pbtxt
@@ -1,15 +1,37 @@
path: "tensorflow.keras.optimizers.Adamax"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Adamax\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adamax.Adamax\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.adam.Adam\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'beta_1\', \'beta_2\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'0.002\', \'0.9\', \'0.999\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'beta_1\', \'beta_2\', \'epsilon\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.999\', \'1e-07\', \'Adamax\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +42,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +58,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-optimizer.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-optimizer.pbtxt
index 53d64da..58b7f27 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-optimizer.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-optimizer.pbtxt
@@ -1,14 +1,35 @@
path: "tensorflow.keras.optimizers.Optimizer"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
+ argspec: "args=[\'self\', \'name\'], varargs=None, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -19,6 +40,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -27,7 +56,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-r-m-sprop.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-r-m-sprop.pbtxt
index a1e9b8c..8de796e 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-r-m-sprop.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-r-m-sprop.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.RMSprop"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.RMSprop\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.rmsprop.RMSprop\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'rho\', \'epsilon\', \'decay\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'None\', \'0.0\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'rho\', \'momentum\', \'epsilon\', \'centered\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.9\', \'0.0\', \'1e-07\', \'False\', \'RMSprop\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-s-g-d.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-s-g-d.pbtxt
index a67fefb..393eeb3 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-s-g-d.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.optimizers.-s-g-d.pbtxt
@@ -1,15 +1,36 @@
path: "tensorflow.keras.optimizers.SGD"
tf_class {
- is_instance: "<class \'tensorflow.python.keras.optimizers.SGD\'>"
- is_instance: "<class \'tensorflow.python.keras.optimizers.Optimizer\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.gradient_descent.SGD\'>"
+ is_instance: "<class \'tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2\'>"
+ is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
is_instance: "<type \'object\'>"
+ member {
+ name: "iterations"
+ mtype: "<type \'property\'>"
+ }
+ member {
+ name: "weights"
+ mtype: "<type \'property\'>"
+ }
member_method {
name: "__init__"
- argspec: "args=[\'self\', \'lr\', \'momentum\', \'decay\', \'nesterov\'], varargs=None, keywords=kwargs, defaults=[\'0.01\', \'0.0\', \'0.0\', \'False\'], "
+ argspec: "args=[\'self\', \'learning_rate\', \'momentum\', \'nesterov\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'0.001\', \'0.0\', \'False\', \'SGD\'], "
+ }
+ member_method {
+ name: "add_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\', \'initializer\'], varargs=None, keywords=None, defaults=[\'zeros\'], "
+ }
+ member_method {
+ name: "add_weight"
+ argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'trainable\', \'synchronization\', \'aggregation\'], varargs=None, keywords=None, defaults=[\'None\', \'zeros\', \'None\', \'VariableSynchronization.AUTO\', \'VariableAggregation.NONE\'], "
+ }
+ member_method {
+ name: "apply_gradients"
+ argspec: "args=[\'self\', \'grads_and_vars\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
+ argspec: "args=[\'cls\', \'config\', \'custom_objects\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "get_config"
@@ -20,6 +41,14 @@
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "get_slot"
+ argspec: "args=[\'self\', \'var\', \'slot_name\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
+ name: "get_slot_names"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
+ member_method {
name: "get_updates"
argspec: "args=[\'self\', \'loss\', \'params\'], varargs=None, keywords=None, defaults=None"
}
@@ -28,7 +57,15 @@
argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
}
member_method {
+ name: "minimize"
+ argspec: "args=[\'self\', \'loss\', \'var_list\', \'grad_loss\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "set_weights"
argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
}
+ member_method {
+ name: "variables"
+ argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+ }
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
index a3599bf..3e1e2e3 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
@@ -133,6 +133,10 @@
argspec: "args=[\'matrix\', \'rhs\', \'l2_regularizer\', \'fast\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'True\', \'None\'], "
}
member_method {
+ name: "lu"
+ argspec: "args=[\'input\', \'output_idx_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'int32\'>\", \'None\'], "
+ }
+ member_method {
name: "matmul"
argspec: "args=[\'a\', \'b\', \'transpose_a\', \'transpose_b\', \'adjoint_a\', \'adjoint_b\', \'a_is_sparse\', \'b_is_sparse\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'False\', \'False\', \'False\', \'False\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt
new file mode 100644
index 0000000..68c651a
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt
@@ -0,0 +1,12 @@
+path: "tensorflow.lite.OpsSet"
+tf_class {
+ is_instance: "<enum \'OpsSet\'>"
+ member {
+ name: "SELECT_TF_OPS"
+ mtype: "<enum \'OpsSet\'>"
+ }
+ member {
+ name: "TFLITE_BUILTINS"
+ mtype: "<enum \'OpsSet\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.lite.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.lite.pbtxt
index f5013c2..154dd00 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.lite.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.lite.pbtxt
@@ -9,6 +9,10 @@
mtype: "<type \'type\'>"
}
member {
+ name: "OpsSet"
+ mtype: "<class \'enum.EnumMeta\'>"
+ }
+ member {
name: "TFLiteConverter"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt
index 979d77e..4ac0484 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.math.pbtxt
@@ -170,7 +170,7 @@
}
member_method {
name: "in_top_k"
- argspec: "args=[\'predictions\', \'targets\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ argspec: "args=[\'targets\', \'predictions\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "invert_permutation"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-accuracy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-accuracy.pbtxt
deleted file mode 100644
index f8e12f8..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.Accuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Accuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'accuracy\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-binary-accuracy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-binary-accuracy.pbtxt
deleted file mode 100644
index b9bc6a7..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-binary-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.BinaryAccuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.BinaryAccuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\', \'threshold\'], varargs=None, keywords=None, defaults=[\'binary_accuracy\', \'None\', \'0.5\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-categorical-accuracy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-categorical-accuracy.pbtxt
deleted file mode 100644
index 0ef75d8..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-categorical-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.CategoricalAccuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.CategoricalAccuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'categorical_accuracy\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-false-negatives.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-false-negatives.pbtxt
deleted file mode 100644
index 33226a2..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-false-negatives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.FalseNegatives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalseNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-false-positives.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-false-positives.pbtxt
deleted file mode 100644
index 9953162..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-false-positives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.FalsePositives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.FalsePositives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-mean.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-mean.pbtxt
deleted file mode 100644
index 7fe6d6f..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-mean.pbtxt
+++ /dev/null
@@ -1,192 +0,0 @@
-path: "tensorflow.metrics.Mean"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'mean\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'values\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-precision.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-precision.pbtxt
deleted file mode 100644
index 8c3271a..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-precision.pbtxt
+++ /dev/null
@@ -1,192 +0,0 @@
-path: "tensorflow.metrics.Precision"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Precision\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-recall.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-recall.pbtxt
deleted file mode 100644
index 840a68b..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-recall.pbtxt
+++ /dev/null
@@ -1,192 +0,0 @@
-path: "tensorflow.metrics.Recall"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.Recall\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-sparse-categorical-accuracy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-sparse-categorical-accuracy.pbtxt
deleted file mode 100644
index 7bce43f..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-sparse-categorical-accuracy.pbtxt
+++ /dev/null
@@ -1,194 +0,0 @@
-path: "tensorflow.metrics.SparseCategoricalAccuracy"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.SparseCategoricalAccuracy\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.MeanMetricWrapper\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Mean\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'sparse_categorical_accuracy\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-true-negatives.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-true-negatives.pbtxt
deleted file mode 100644
index 83cd5b7..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-true-negatives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.TrueNegatives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.TrueNegatives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-true-positives.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.-true-positives.pbtxt
deleted file mode 100644
index 5b2502e..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.-true-positives.pbtxt
+++ /dev/null
@@ -1,193 +0,0 @@
-path: "tensorflow.metrics.TruePositives"
-tf_class {
- is_instance: "<class \'tensorflow.python.keras.metrics.TruePositives\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics._ConfusionMatrixConditionCount\'>"
- is_instance: "<class \'tensorflow.python.keras.metrics.Metric\'>"
- is_instance: "<class \'tensorflow.python.keras.engine.base_layer.Layer\'>"
- is_instance: "<class \'tensorflow.python.training.checkpointable.base.CheckpointableBase\'>"
- is_instance: "<type \'object\'>"
- member {
- name: "activity_regularizer"
- mtype: "<type \'property\'>"
- }
- member {
- name: "dtype"
- mtype: "<type \'property\'>"
- }
- member {
- name: "inbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "input_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "losses"
- mtype: "<type \'property\'>"
- }
- member {
- name: "name"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "non_trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "outbound_nodes"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_mask"
- mtype: "<type \'property\'>"
- }
- member {
- name: "output_shape"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "trainable_weights"
- mtype: "<type \'property\'>"
- }
- member {
- name: "updates"
- mtype: "<type \'property\'>"
- }
- member {
- name: "variables"
- mtype: "<type \'property\'>"
- }
- member {
- name: "weights"
- mtype: "<type \'property\'>"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\', \'thresholds\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
- }
- member_method {
- name: "add_loss"
- argspec: "args=[\'self\', \'losses\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_metric"
- argspec: "args=[\'self\', \'value\', \'aggregation\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
- }
- member_method {
- name: "add_update"
- argspec: "args=[\'self\', \'updates\', \'inputs\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "add_variable"
- argspec: "args=[\'self\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "add_weight"
- argspec: "args=[\'self\', \'name\', \'shape\', \'aggregation\', \'synchronization\', \'initializer\'], varargs=None, keywords=None, defaults=[\'()\', \'VariableAggregation.SUM\', \'VariableSynchronization.ON_READ\', \'None\'], "
- }
- member_method {
- name: "apply"
- argspec: "args=[\'self\', \'inputs\'], varargs=args, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "build"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "call"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None"
- }
- member_method {
- name: "compute_mask"
- argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
- name: "compute_output_shape"
- argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "count_params"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "from_config"
- argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_config"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_input_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_losses_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_mask_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_output_shape_at"
- argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_updates_for"
- argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "get_weights"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "reset_states"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "result"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "set_weights"
- argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "update_state"
- argspec: "args=[\'self\', \'y_true\', \'y_pred\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.metrics.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.metrics.pbtxt
deleted file mode 100644
index 773efd0..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.metrics.pbtxt
+++ /dev/null
@@ -1,47 +0,0 @@
-path: "tensorflow.metrics"
-tf_module {
- member {
- name: "Accuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "BinaryAccuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "CategoricalAccuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "FalseNegatives"
- mtype: "<type \'type\'>"
- }
- member {
- name: "FalsePositives"
- mtype: "<type \'type\'>"
- }
- member {
- name: "Mean"
- mtype: "<type \'type\'>"
- }
- member {
- name: "Precision"
- mtype: "<type \'type\'>"
- }
- member {
- name: "Recall"
- mtype: "<type \'type\'>"
- }
- member {
- name: "SparseCategoricalAccuracy"
- mtype: "<type \'type\'>"
- }
- member {
- name: "TrueNegatives"
- mtype: "<type \'type\'>"
- }
- member {
- name: "TruePositives"
- mtype: "<type \'type\'>"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
index 63bf24b..c75c75f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
@@ -106,7 +106,7 @@
}
member_method {
name: "depth_to_space"
- argspec: "args=[\'input\', \'block_size\', \'name\', \'data_format\'], varargs=None, keywords=None, defaults=[\'None\', \'NHWC\'], "
+ argspec: "args=[\'input\', \'block_size\', \'data_format\', \'name\'], varargs=None, keywords=None, defaults=[\'NHWC\', \'None\'], "
}
member_method {
name: "depthwise_conv2d"
@@ -134,11 +134,11 @@
}
member_method {
name: "embedding_lookup"
- argspec: "args=[\'params\', \'ids\', \'partition_strategy\', \'name\', \'validate_indices\', \'max_norm\'], varargs=None, keywords=None, defaults=[\'mod\', \'None\', \'True\', \'None\'], "
+ argspec: "args=[\'params\', \'ids\', \'partition_strategy\', \'max_norm\', \'name\'], varargs=None, keywords=None, defaults=[\'mod\', \'None\', \'None\'], "
}
member_method {
name: "embedding_lookup_sparse"
- argspec: "args=[\'params\', \'sp_ids\', \'sp_weights\', \'partition_strategy\', \'name\', \'combiner\', \'max_norm\'], varargs=None, keywords=None, defaults=[\'mod\', \'None\', \'None\', \'None\'], "
+ argspec: "args=[\'params\', \'sp_ids\', \'sp_weights\', \'partition_strategy\', \'combiner\', \'max_norm\', \'name\'], varargs=None, keywords=None, defaults=[\'mod\', \'None\', \'None\', \'None\'], "
}
member_method {
name: "erosion2d"
@@ -158,7 +158,7 @@
}
member_method {
name: "in_top_k"
- argspec: "args=[\'predictions\', \'targets\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ argspec: "args=[\'targets\', \'predictions\', \'k\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "l2_loss"
@@ -229,16 +229,12 @@
argspec: "args=[\'features\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "relu_layer"
- argspec: "args=[\'x\', \'weights\', \'biases\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
- }
- member_method {
name: "safe_embedding_lookup_sparse"
- argspec: "args=[\'embedding_weights\', \'sparse_ids\', \'sparse_weights\', \'combiner\', \'default_id\', \'name\', \'partition_strategy\', \'max_norm\'], varargs=None, keywords=None, defaults=[\'None\', \'mean\', \'None\', \'None\', \'div\', \'None\'], "
+ argspec: "args=[\'embedding_weights\', \'sparse_ids\', \'sparse_weights\', \'combiner\', \'default_id\', \'max_norm\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'mean\', \'None\', \'None\', \'None\'], "
}
member_method {
name: "sampled_softmax_loss"
- argspec: "args=[\'weights\', \'biases\', \'labels\', \'inputs\', \'num_sampled\', \'num_classes\', \'num_true\', \'sampled_values\', \'remove_accidental_hits\', \'partition_strategy\', \'name\', \'seed\'], varargs=None, keywords=None, defaults=[\'1\', \'None\', \'True\', \'mod\', \'sampled_softmax_loss\', \'None\'], "
+ argspec: "args=[\'weights\', \'biases\', \'labels\', \'inputs\', \'num_sampled\', \'num_classes\', \'num_true\', \'sampled_values\', \'remove_accidental_hits\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'None\', \'True\', \'None\', \'sampled_softmax_loss\'], "
}
member_method {
name: "selu"
@@ -274,11 +270,11 @@
}
member_method {
name: "space_to_batch"
- argspec: "args=[\'input\', \'paddings\', \'block_size\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ argspec: "args=[\'input\', \'block_shape\', \'paddings\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "space_to_depth"
- argspec: "args=[\'input\', \'block_size\', \'name\', \'data_format\'], varargs=None, keywords=None, defaults=[\'None\', \'NHWC\'], "
+ argspec: "args=[\'input\', \'block_size\', \'data_format\', \'name\'], varargs=None, keywords=None, defaults=[\'NHWC\', \'None\'], "
}
member_method {
name: "sparse_softmax_cross_entropy_with_logits"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.pbtxt
index cb38ae0..574b677 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.pbtxt
@@ -5,14 +5,6 @@
mtype: "<type \'type\'>"
}
member {
- name: "AttrValue"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "ConfigProto"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
name: "DType"
mtype: "<type \'type\'>"
}
@@ -21,14 +13,6 @@
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
member {
- name: "FIFOQueue"
- mtype: "<type \'type\'>"
- }
- member {
- name: "GPUOptions"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
name: "GradientTape"
mtype: "<type \'type\'>"
}
@@ -37,70 +21,26 @@
mtype: "<type \'type\'>"
}
member {
- name: "GraphDef"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "GraphOptions"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "HistogramProto"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
name: "IndexedSlices"
mtype: "<type \'type\'>"
}
member {
- name: "LogMessage"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "MetaGraphDef"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "NameAttrList"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "NodeDef"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
name: "Operation"
mtype: "<type \'type\'>"
}
member {
- name: "OptimizerOptions"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+ name: "RaggedTensor"
+ mtype: "<type \'type\'>"
}
member {
name: "RegisterGradient"
mtype: "<type \'type\'>"
}
member {
- name: "RunMetadata"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "RunOptions"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
- name: "SessionLog"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
name: "SparseTensor"
mtype: "<type \'type\'>"
}
member {
- name: "SparseTensorValue"
- mtype: "<type \'type\'>"
- }
- member {
name: "Summary"
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
@@ -141,6 +81,10 @@
mtype: "<class \'enum.EnumMeta\'>"
}
member {
+ name: "autograph"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "bfloat16"
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
@@ -277,10 +221,6 @@
mtype: "<type \'module\'>"
}
member {
- name: "metrics"
- mtype: "<type \'module\'>"
- }
- member {
name: "name_scope"
mtype: "<type \'type\'>"
}
@@ -313,6 +253,10 @@
mtype: "<type \'module\'>"
}
member {
+ name: "queue"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "quint16"
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
@@ -321,6 +265,10 @@
mtype: "<class \'tensorflow.python.framework.dtypes.DType\'>"
}
member {
+ name: "ragged"
+ mtype: "<type \'module\'>"
+ }
+ member {
name: "random"
mtype: "<type \'module\'>"
}
@@ -493,10 +441,6 @@
argspec: "args=[\'params\', \'indices\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
- name: "batch_scatter_update"
- argspec: "args=[\'ref\', \'indices\', \'updates\', \'use_locking\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
- }
- member_method {
name: "batch_to_space"
argspec: "args=[\'input\', \'block_shape\', \'crops\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -997,6 +941,10 @@
argspec: "args=[\'values\', \'axis\', \'direction\', \'name\'], varargs=None, keywords=None, defaults=[\'-1\', \'ASCENDING\', \'None\'], "
}
member_method {
+ name: "space_to_batch"
+ argspec: "args=[\'input\', \'block_shape\', \'paddings\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "space_to_batch_nd"
argspec: "args=[\'input\', \'block_shape\', \'paddings\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
@@ -1033,10 +981,6 @@
argspec: "args=[\'source\', \'delimiter\', \'skip_empty\'], varargs=None, keywords=None, defaults=[\' \', \'True\'], "
}
member_method {
- name: "substr"
- argspec: "args=[\'input\', \'pos\', \'len\', \'name\', \'unit\'], varargs=None, keywords=None, defaults=[\'None\', \'BYTE\'], "
- }
- member_method {
name: "subtract"
argspec: "args=[\'x\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.-f-i-f-o-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.queue.-f-i-f-o-queue.pbtxt
similarity index 98%
rename from tensorflow/tools/api/golden/v2/tensorflow.-f-i-f-o-queue.pbtxt
rename to tensorflow/tools/api/golden/v2/tensorflow.queue.-f-i-f-o-queue.pbtxt
index a095616..724ab5f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.-f-i-f-o-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.queue.-f-i-f-o-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.FIFOQueue"
+path: "tensorflow.queue.FIFOQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.FIFOQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.queue.-padding-f-i-f-o-queue.pbtxt
similarity index 97%
copy from tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.queue.-padding-f-i-f-o-queue.pbtxt
index 85306fd..9ef0a4d 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-padding-f-i-f-o-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.queue.-padding-f-i-f-o-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.PaddingFIFOQueue"
+path: "tensorflow.queue.PaddingFIFOQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PaddingFIFOQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.queue.-priority-queue.pbtxt
similarity index 97%
copy from tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.queue.-priority-queue.pbtxt
index 02d8037..bb66beb 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-priority-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.queue.-priority-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.PriorityQueue"
+path: "tensorflow.queue.PriorityQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.PriorityQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.queue.-queue-base.pbtxt
similarity index 97%
copy from tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.queue.-queue-base.pbtxt
index a30481a..8faaad2 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-queue-base.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.queue.-queue-base.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.QueueBase"
+path: "tensorflow.queue.QueueBase"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
is_instance: "<type \'object\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.queue.-random-shuffle-queue.pbtxt
similarity index 97%
copy from tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt
copy to tensorflow/tools/api/golden/v2/tensorflow.queue.-random-shuffle-queue.pbtxt
index 82cbf98..31cd503 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.io.-random-shuffle-queue.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.queue.-random-shuffle-queue.pbtxt
@@ -1,4 +1,4 @@
-path: "tensorflow.io.RandomShuffleQueue"
+path: "tensorflow.queue.RandomShuffleQueue"
tf_class {
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.RandomShuffleQueue\'>"
is_instance: "<class \'tensorflow.python.ops.data_flow_ops.QueueBase\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.queue.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.queue.pbtxt
new file mode 100644
index 0000000..c16e95e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.queue.pbtxt
@@ -0,0 +1,23 @@
+path: "tensorflow.queue"
+tf_module {
+ member {
+ name: "FIFOQueue"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "PaddingFIFOQueue"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "PriorityQueue"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "QueueBase"
+ mtype: "<type \'type\'>"
+ }
+ member {
+ name: "RandomShuffleQueue"
+ mtype: "<type \'type\'>"
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.ragged.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.ragged.pbtxt
new file mode 100644
index 0000000..5fde488
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.ragged.pbtxt
@@ -0,0 +1,23 @@
+path: "tensorflow.ragged"
+tf_module {
+ member_method {
+ name: "constant"
+ argspec: "args=[\'pylist\', \'dtype\', \'ragged_rank\', \'inner_shape\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "map_flat_values"
+ argspec: "args=[\'op\'], varargs=args, keywords=kwargs, defaults=None"
+ }
+ member_method {
+ name: "range"
+ argspec: "args=[\'starts\', \'limits\', \'deltas\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'None\', \'None\'], "
+ }
+ member_method {
+ name: "row_splits_to_segment_ids"
+ argspec: "args=[\'splits\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "segment_ids_to_row_splits"
+ argspec: "args=[\'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
index de5cb6b..d49c23e 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
@@ -17,6 +17,10 @@
argspec: "args=[\'shape\', \'alpha\', \'beta\', \'dtype\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \"<dtype: \'float32\'>\", \'None\', \'None\'], "
}
member_method {
+ name: "learned_unigram_candidate_sampler"
+ argspec: "args=[\'true_classes\', \'num_true\', \'num_sampled\', \'unique\', \'range_max\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
+ }
+ member_method {
name: "log_uniform_candidate_sampler"
argspec: "args=[\'true_classes\', \'num_true\', \'num_sampled\', \'unique\', \'range_max\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt
index 9808200..b8bd2c0 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.sparse.pbtxt
@@ -14,7 +14,7 @@
}
member_method {
name: "concat"
- argspec: "args=[\'axis\', \'sp_inputs\', \'expand_nonconcat_dim\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
+ argspec: "args=[\'axis\', \'sp_inputs\', \'expand_nonconcat_dims\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
}
member_method {
name: "cross"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt
index f6e32ed..962cf9a 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.strings.pbtxt
@@ -53,6 +53,14 @@
argspec: "args=[\'input\', \'out_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'float32\'>\", \'None\'], "
}
member_method {
+ name: "unicode_decode"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'replace_control_characters\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'False\', \'None\'], "
+ }
+ member_method {
+ name: "unicode_decode_with_offsets"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'replace_control_characters\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'False\', \'None\'], "
+ }
+ member_method {
name: "unicode_encode"
argspec: "args=[\'input\', \'output_encoding\', \'errors\', \'replacement_char\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'None\'], "
}
@@ -61,6 +69,14 @@
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unicode_split"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'None\'], "
+ }
+ member_method {
+ name: "unicode_split_with_offsets"
+ argspec: "args=[\'input\', \'input_encoding\', \'errors\', \'replacement_char\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'None\'], "
+ }
+ member_method {
name: "unicode_transcode"
argspec: "args=[\'input\', \'input_encoding\', \'output_encoding\', \'errors\', \'replacement_char\', \'replace_control_characters\', \'name\'], varargs=None, keywords=None, defaults=[\'replace\', \'65533\', \'False\', \'None\'], "
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.summary.-session-log.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.summary.-session-log.pbtxt
deleted file mode 100644
index 73de738..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.summary.-session-log.pbtxt
+++ /dev/null
@@ -1,44 +0,0 @@
-path: "tensorflow.summary.SessionLog"
-tf_proto {
- descriptor {
- name: "SessionLog"
- field {
- name: "status"
- number: 1
- label: LABEL_OPTIONAL
- type: TYPE_ENUM
- type_name: ".tensorflow.SessionLog.SessionStatus"
- }
- field {
- name: "checkpoint_path"
- number: 2
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- field {
- name: "msg"
- number: 3
- label: LABEL_OPTIONAL
- type: TYPE_STRING
- }
- enum_type {
- name: "SessionStatus"
- value {
- name: "STATUS_UNSPECIFIED"
- number: 0
- }
- value {
- name: "START"
- number: 1
- }
- value {
- name: "STOP"
- number: 2
- }
- value {
- name: "CHECKPOINT"
- number: 3
- }
- }
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.summary.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.summary.pbtxt
index 42a74a6..5cf4d7c 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.summary.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.summary.pbtxt
@@ -13,10 +13,6 @@
mtype: "<type \'type\'>"
}
member {
- name: "SessionLog"
- mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
- }
- member {
name: "Summary"
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.test.-stub-out-for-testing.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.test.-stub-out-for-testing.pbtxt
deleted file mode 100644
index e02a0c6..0000000
--- a/tensorflow/tools/api/golden/v2/tensorflow.test.-stub-out-for-testing.pbtxt
+++ /dev/null
@@ -1,28 +0,0 @@
-path: "tensorflow.test.StubOutForTesting"
-tf_class {
- is_instance: "<class \'tensorflow.python.platform.googletest.StubOutForTesting\'>"
- member_method {
- name: "CleanUp"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "Set"
- argspec: "args=[\'self\', \'parent\', \'child_name\', \'new_child\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "SmartSet"
- argspec: "args=[\'self\', \'obj\', \'attr_name\', \'new_attr\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "SmartUnsetAll"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "UnsetAll"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
- member_method {
- name: "__init__"
- argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
- }
-}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt
index 72ce733..980e96a 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.test.pbtxt
@@ -5,10 +5,6 @@
mtype: "<class \'tensorflow.python.platform.benchmark._BenchmarkRegistrar\'>"
}
member {
- name: "StubOutForTesting"
- mtype: "<type \'type\'>"
- }
- member {
name: "TestCase"
mtype: "<type \'type\'>"
}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt
index 3ff4b69..c72564e 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.train.pbtxt
@@ -13,10 +13,6 @@
mtype: "<type \'type\'>"
}
member {
- name: "CheckpointSaverHook"
- mtype: "<type \'type\'>"
- }
- member {
name: "ClusterDef"
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
@@ -53,22 +49,10 @@
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
member {
- name: "FeedFnHook"
- mtype: "<type \'type\'>"
- }
- member {
- name: "FinalOpsHook"
- mtype: "<type \'type\'>"
- }
- member {
name: "FloatList"
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
member {
- name: "GlobalStepWaiterHook"
- mtype: "<type \'type\'>"
- }
- member {
name: "Int64List"
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
@@ -77,14 +61,6 @@
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
member {
- name: "LoggingTensorHook"
- mtype: "<type \'type\'>"
- }
- member {
- name: "NanTensorHook"
- mtype: "<type \'type\'>"
- }
- member {
name: "ProximalGradientDescentOptimizer"
mtype: "<type \'type\'>"
}
@@ -93,29 +69,9 @@
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
member {
- name: "Server"
- mtype: "<type \'type\'>"
- }
- member {
name: "ServerDef"
mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
}
- member {
- name: "SessionRunHook"
- mtype: "<type \'type\'>"
- }
- member {
- name: "StepCounterHook"
- mtype: "<type \'type\'>"
- }
- member {
- name: "StopAtStepHook"
- mtype: "<type \'type\'>"
- }
- member {
- name: "SummarySaverHook"
- mtype: "<type \'type\'>"
- }
member_method {
name: "cosine_decay"
argspec: "args=[\'learning_rate\', \'global_step\', \'decay_steps\', \'alpha\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'None\'], "
diff --git a/tensorflow/tools/api/lib/python_object_to_proto_visitor.py b/tensorflow/tools/api/lib/python_object_to_proto_visitor.py
index 70df38b..5102066 100644
--- a/tensorflow/tools/api/lib/python_object_to_proto_visitor.py
+++ b/tensorflow/tools/api/lib/python_object_to_proto_visitor.py
@@ -37,6 +37,9 @@
'train.NanLossDuringTrainingError': {
'message': {}
},
+ 'estimator.NanLossDuringTrainingError': {
+ 'message': {}
+ },
}
# Python 2 vs. 3 differences
diff --git a/tensorflow/tools/api/tests/api_compatibility_test.py b/tensorflow/tools/api/tests/api_compatibility_test.py
index e7f23a1..723fcee 100644
--- a/tensorflow/tools/api/tests/api_compatibility_test.py
+++ b/tensorflow/tools/api/tests/api_compatibility_test.py
@@ -277,6 +277,9 @@
public_api_visitor = public_api.PublicAPIVisitor(visitor)
public_api_visitor.private_map['tf'] = ['contrib']
+ if api_version == 2:
+ public_api_visitor.private_map['tf'].append('enable_v2_behavior')
+
public_api_visitor.do_not_descend_map['tf.GPUOptions'] = ['Experimental']
if FLAGS.only_test_core_api:
public_api_visitor.do_not_descend_map['tf'].extend(_NON_CORE_PACKAGES)
@@ -311,7 +314,7 @@
update_goldens=FLAGS.update_goldens,
api_version=api_version)
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testAPIBackwardsCompatibility(self):
api_version = 1
golden_file_pattern = os.path.join(
@@ -330,7 +333,7 @@
'tensorflow.python.util.lazy_loader.LazyLoader'
in str(type(tf.contrib)))
- @test_util.run_deprecated_v1
+ @test_util.run_v1_only('b/120545219')
def testAPIBackwardsCompatibilityV1(self):
api_version = 1
golden_file_pattern = os.path.join(
diff --git a/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh b/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh
index a6871dd..4c4e8ba 100644
--- a/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh
+++ b/tensorflow/tools/ci_build/windows/cpu/pip/build_tf_windows.sh
@@ -109,7 +109,7 @@
if [ -z ${PROJECT_NAME} ]; then
EXTRA_PIP_FLAGS="--nightly_flag"
else
- EXTRA_PIP_FLAGS="--project_name=${PROJECT_NAME} --nightly_flag"
+ EXTRA_PIP_FLAGS="--project_name ${PROJECT_NAME} --nightly_flag"
fi
fi
@@ -123,7 +123,8 @@
run_configure_for_cpu_build
bazel build --announce_rc --config=opt ${EXTRA_BUILD_FLAGS} \
- tensorflow/tools/pip_package:build_pip_package || exit $?
+ tensorflow/tools/pip_package:build_pip_package \
+ --incompatible_remove_native_http_archive=false || exit $?
if [[ "$SKIP_TEST" == 1 ]]; then
exit 0
@@ -132,7 +133,7 @@
# Create a python test directory to avoid package name conflict
create_python_test_dir "${PY_TEST_DIR}"
-./bazel-bin/tensorflow/tools/pip_package/build_pip_package "$PWD/${PY_TEST_DIR}" "${EXTRA_PIP_FLAGS}"
+./bazel-bin/tensorflow/tools/pip_package/build_pip_package "$PWD/${PY_TEST_DIR}" ${EXTRA_PIP_FLAGS}
if [[ "$TF_NIGHTLY" == 1 ]]; then
exit 0
diff --git a/tensorflow/tools/ci_build/windows/gpu/pip/build_tf_windows.sh b/tensorflow/tools/ci_build/windows/gpu/pip/build_tf_windows.sh
index acafd9e..070235f 100644
--- a/tensorflow/tools/ci_build/windows/gpu/pip/build_tf_windows.sh
+++ b/tensorflow/tools/ci_build/windows/gpu/pip/build_tf_windows.sh
@@ -136,7 +136,7 @@
create_python_test_dir "${PY_TEST_DIR}"
./bazel-bin/tensorflow/tools/pip_package/build_pip_package "$PWD/${PY_TEST_DIR}" \
- --gpu "${EXTRA_PIP_FLAGS}"
+ --gpu ${EXTRA_PIP_FLAGS}
if [[ "$TF_NIGHTLY" == 1 ]]; then
exit 0
diff --git a/tensorflow/tools/compatibility/BUILD b/tensorflow/tools/compatibility/BUILD
index 05d924c..a9902d7 100644
--- a/tensorflow/tools/compatibility/BUILD
+++ b/tensorflow/tools/compatibility/BUILD
@@ -52,13 +52,20 @@
)
py_library(
- name = "tf_upgrade_v2_lib",
- srcs = [
- "renames_v2.py",
- "tf_upgrade_v2.py",
- ],
+ name = "reorders_v2",
+ srcs = ["reorders_v2.py"],
srcs_version = "PY2AND3",
- deps = [":ast_edits"],
+)
+
+py_library(
+ name = "tf_upgrade_v2_lib",
+ srcs = ["tf_upgrade_v2.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ ":ast_edits",
+ ":renames_v2",
+ ":reorders_v2",
+ ],
)
py_binary(
diff --git a/tensorflow/tools/compatibility/README.md b/tensorflow/tools/compatibility/README.md
index 6ff42b1..5e2de35 100644
--- a/tensorflow/tools/compatibility/README.md
+++ b/tensorflow/tools/compatibility/README.md
@@ -7,7 +7,7 @@
## Running the script from pip package
-First, install TensorFlow pip package. See
+First, install TensorFlow pip package*. See
https://www.tensorflow.org/install/pip.
Upgrade script can be run on a single Python file:
@@ -27,6 +27,8 @@
tf_upgrade_v2 --intree coolcode --outtree coolcode-upgraded --copyotherfiles False
```
+*Note: `tf_upgrade_v2` is installed automatically as a script by the pip install
+after TensorFlow 1.12.
## Report
diff --git a/tensorflow/tools/compatibility/renames_v2.py b/tensorflow/tools/compatibility/renames_v2.py
index 7fe3d1e..2763a0c 100644
--- a/tensorflow/tools/compatibility/renames_v2.py
+++ b/tensorflow/tools/compatibility/renames_v2.py
@@ -26,37 +26,53 @@
renames = {
'tf.AUTO_REUSE': 'tf.compat.v1.AUTO_REUSE',
+ 'tf.AttrValue': 'tf.compat.v1.AttrValue',
'tf.COMPILER_VERSION': 'tf.version.COMPILER_VERSION',
'tf.CXX11_ABI_FLAG': 'tf.sysconfig.CXX11_ABI_FLAG',
'tf.ConditionalAccumulator': 'tf.compat.v1.ConditionalAccumulator',
'tf.ConditionalAccumulatorBase': 'tf.compat.v1.ConditionalAccumulatorBase',
+ 'tf.ConfigProto': 'tf.compat.v1.ConfigProto',
'tf.DeviceSpec': 'tf.compat.v1.DeviceSpec',
'tf.Dimension': 'tf.compat.v1.Dimension',
+ 'tf.FIFOQueue': 'tf.queue.FIFOQueue',
'tf.FixedLenFeature': 'tf.io.FixedLenFeature',
'tf.FixedLenSequenceFeature': 'tf.io.FixedLenSequenceFeature',
'tf.FixedLengthRecordReader': 'tf.compat.v1.FixedLengthRecordReader',
'tf.GIT_VERSION': 'tf.version.GIT_VERSION',
+ 'tf.GPUOptions': 'tf.compat.v1.GPUOptions',
'tf.GRAPH_DEF_VERSION': 'tf.version.GRAPH_DEF_VERSION',
'tf.GRAPH_DEF_VERSION_MIN_CONSUMER': 'tf.version.GRAPH_DEF_VERSION_MIN_CONSUMER',
'tf.GRAPH_DEF_VERSION_MIN_PRODUCER': 'tf.version.GRAPH_DEF_VERSION_MIN_PRODUCER',
+ 'tf.GraphDef': 'tf.compat.v1.GraphDef',
'tf.GraphKeys': 'tf.compat.v1.GraphKeys',
+ 'tf.GraphOptions': 'tf.compat.v1.GraphOptions',
+ 'tf.HistogramProto': 'tf.compat.v1.HistogramProto',
'tf.IdentityReader': 'tf.compat.v1.IdentityReader',
'tf.InteractiveSession': 'tf.compat.v1.InteractiveSession',
'tf.LMDBReader': 'tf.compat.v1.LMDBReader',
+ 'tf.LogMessage': 'tf.compat.v1.LogMessage',
'tf.MONOLITHIC_BUILD': 'tf.sysconfig.MONOLITHIC_BUILD',
+ 'tf.MetaGraphDef': 'tf.compat.v1.MetaGraphDef',
+ 'tf.NameAttrList': 'tf.compat.v1.NameAttrList',
'tf.NoGradient': 'tf.no_gradient',
+ 'tf.NodeDef': 'tf.compat.v1.NodeDef',
'tf.NotDifferentiable': 'tf.no_gradient',
'tf.OpError': 'tf.errors.OpError',
- 'tf.PaddingFIFOQueue': 'tf.io.PaddingFIFOQueue',
+ 'tf.OptimizerOptions': 'tf.compat.v1.OptimizerOptions',
+ 'tf.PaddingFIFOQueue': 'tf.queue.PaddingFIFOQueue',
'tf.Print': 'tf.compat.v1.Print',
- 'tf.PriorityQueue': 'tf.io.PriorityQueue',
+ 'tf.PriorityQueue': 'tf.queue.PriorityQueue',
'tf.QUANTIZED_DTYPES': 'tf.dtypes.QUANTIZED_DTYPES',
- 'tf.QueueBase': 'tf.io.QueueBase',
- 'tf.RandomShuffleQueue': 'tf.io.RandomShuffleQueue',
+ 'tf.QueueBase': 'tf.queue.QueueBase',
+ 'tf.RandomShuffleQueue': 'tf.queue.RandomShuffleQueue',
'tf.ReaderBase': 'tf.compat.v1.ReaderBase',
+ 'tf.RunMetadata': 'tf.compat.v1.RunMetadata',
+ 'tf.RunOptions': 'tf.compat.v1.RunOptions',
'tf.Session': 'tf.compat.v1.Session',
+ 'tf.SessionLog': 'tf.compat.v1.SessionLog',
'tf.SparseConditionalAccumulator': 'tf.sparse.SparseConditionalAccumulator',
'tf.SparseFeature': 'tf.io.SparseFeature',
+ 'tf.SparseTensorValue': 'tf.compat.v1.SparseTensorValue',
'tf.TFRecordReader': 'tf.compat.v1.TFRecordReader',
'tf.TensorInfo': 'tf.compat.v1.TensorInfo',
'tf.TextLineReader': 'tf.compat.v1.TextLineReader',
@@ -90,6 +106,7 @@
'tf.assign': 'tf.compat.v1.assign',
'tf.assign_add': 'tf.compat.v1.assign_add',
'tf.assign_sub': 'tf.compat.v1.assign_sub',
+ 'tf.batch_scatter_update': 'tf.compat.v1.batch_scatter_update',
'tf.betainc': 'tf.math.betainc',
'tf.ceil': 'tf.math.ceil',
'tf.check_numerics': 'tf.debugging.check_numerics',
@@ -118,7 +135,7 @@
'tf.decode_json_example': 'tf.io.decode_json_example',
'tf.decode_raw': 'tf.io.decode_raw',
'tf.delete_session_tensor': 'tf.compat.v1.delete_session_tensor',
- 'tf.depth_to_space': 'tf.nn.depth_to_space',
+ 'tf.depth_to_space': 'tf.compat.v1.depth_to_space',
'tf.dequantize': 'tf.quantization.dequantize',
'tf.deserialize_many_sparse': 'tf.io.deserialize_many_sparse',
'tf.diag': 'tf.linalg.tensor_diag',
@@ -213,6 +230,10 @@
'tf.initializers.tables_initializer': 'tf.compat.v1.initializers.tables_initializer',
'tf.initializers.variables': 'tf.compat.v1.initializers.variables',
'tf.invert_permutation': 'tf.math.invert_permutation',
+ 'tf.io.PaddingFIFOQueue': 'tf.queue.PaddingFIFOQueue',
+ 'tf.io.PriorityQueue': 'tf.queue.PriorityQueue',
+ 'tf.io.QueueBase': 'tf.queue.QueueBase',
+ 'tf.io.RandomShuffleQueue': 'tf.queue.RandomShuffleQueue',
'tf.io.tf_record_iterator': 'tf.compat.v1.io.tf_record_iterator',
'tf.is_finite': 'tf.math.is_finite',
'tf.is_inf': 'tf.math.is_inf',
@@ -222,6 +243,7 @@
'tf.is_strictly_increasing': 'tf.math.is_strictly_increasing',
'tf.is_variable_initialized': 'tf.compat.v1.is_variable_initialized',
'tf.keras.backend.get_session': 'tf.compat.v1.keras.backend.get_session',
+ 'tf.keras.layers.CuDNNLSTM': 'tf.compat.v1.keras.layers.CuDNNLSTM',
'tf.layers.AveragePooling1D': 'tf.compat.v1.layers.AveragePooling1D',
'tf.layers.AveragePooling2D': 'tf.compat.v1.layers.AveragePooling2D',
'tf.layers.AveragePooling3D': 'tf.compat.v1.layers.AveragePooling3D',
@@ -370,6 +392,7 @@
'tf.nn.quantized_max_pool': 'tf.compat.v1.nn.quantized_max_pool',
'tf.nn.quantized_relu_x': 'tf.compat.v1.nn.quantized_relu_x',
'tf.nn.raw_rnn': 'tf.compat.v1.nn.raw_rnn',
+ 'tf.nn.relu_layer': 'tf.compat.v1.nn.relu_layer',
'tf.nn.rnn_cell.BasicLSTMCell': 'tf.compat.v1.nn.rnn_cell.BasicLSTMCell',
'tf.nn.rnn_cell.BasicRNNCell': 'tf.compat.v1.nn.rnn_cell.BasicRNNCell',
'tf.nn.rnn_cell.GRUCell': 'tf.compat.v1.nn.rnn_cell.GRUCell',
@@ -403,6 +426,9 @@
'tf.qr': 'tf.linalg.qr',
'tf.quantize': 'tf.quantization.quantize',
'tf.quantized_concat': 'tf.quantization.quantized_concat',
+ 'tf.ragged.constant_value': 'tf.compat.v1.ragged.constant_value',
+ 'tf.ragged.convert_to_tensor_or_ragged_tensor': 'tf.compat.v1.ragged.convert_to_tensor_or_ragged_tensor',
+ 'tf.ragged.RaggedTensorValue': 'tf.compat.v1.ragged.RaggedTensorValue',
'tf.random.get_seed': 'tf.compat.v1.random.get_seed',
'tf.random.set_random_seed': 'tf.compat.v1.random.set_random_seed',
'tf.random_crop': 'tf.image.random_crop',
@@ -494,8 +520,7 @@
'tf.sets.set_intersection': 'tf.sets.intersection',
'tf.sets.set_size': 'tf.sets.size',
'tf.sets.set_union': 'tf.sets.union',
- 'tf.space_to_batch': 'tf.nn.space_to_batch',
- 'tf.space_to_depth': 'tf.nn.space_to_depth',
+ 'tf.space_to_depth': 'tf.compat.v1.space_to_depth',
'tf.sparse.matmul': 'tf.sparse.sparse_dense_matmul',
'tf.sparse.merge': 'tf.compat.v1.sparse.merge',
'tf.sparse.placeholder': 'tf.compat.v1.sparse.placeholder',
@@ -507,9 +532,7 @@
'tf.sparse_merge': 'tf.compat.v1.sparse_merge',
'tf.sparse_minimum': 'tf.sparse.minimum',
'tf.sparse_placeholder': 'tf.compat.v1.sparse_placeholder',
- 'tf.sparse_reduce_max': 'tf.compat.v1.sparse_reduce_max',
'tf.sparse_reduce_max_sparse': 'tf.compat.v1.sparse_reduce_max_sparse',
- 'tf.sparse_reduce_sum': 'tf.compat.v1.sparse_reduce_sum',
'tf.sparse_reduce_sum_sparse': 'tf.compat.v1.sparse_reduce_sum_sparse',
'tf.sparse_reorder': 'tf.sparse.reorder',
'tf.sparse_reset_shape': 'tf.sparse.reset_shape',
@@ -544,6 +567,7 @@
'tf.string_strip': 'tf.strings.strip',
'tf.string_to_hash_bucket_fast': 'tf.strings.to_hash_bucket_fast',
'tf.string_to_hash_bucket_strong': 'tf.strings.to_hash_bucket_strong',
+ 'tf.summary.SessionLog': 'tf.compat.v1.summary.SessionLog',
'tf.summary.audio': 'tf.compat.v1.summary.audio',
'tf.summary.get_summary_description': 'tf.compat.v1.summary.get_summary_description',
'tf.summary.histogram': 'tf.compat.v1.summary.histogram',
@@ -555,6 +579,7 @@
'tf.summary.text': 'tf.compat.v1.summary.text',
'tf.svd': 'tf.linalg.svd',
'tf.tables_initializer': 'tf.compat.v1.tables_initializer',
+ 'tf.test.StubOutForTesting': 'tf.compat.v1.test.StubOutForTesting',
'tf.test.compute_gradient': 'tf.compat.v1.test.compute_gradient',
'tf.test.compute_gradient_error': 'tf.compat.v1.test.compute_gradient_error',
'tf.test.get_temp_dir': 'tf.compat.v1.test.get_temp_dir',
@@ -572,31 +597,42 @@
'tf.train.AdagradDAOptimizer': 'tf.compat.v1.train.AdagradDAOptimizer',
'tf.train.AdagradOptimizer': 'tf.compat.v1.train.AdagradOptimizer',
'tf.train.AdamOptimizer': 'tf.compat.v1.train.AdamOptimizer',
- 'tf.train.CheckpointSaverListener': 'tf.compat.v1.train.CheckpointSaverListener',
+ 'tf.train.CheckpointSaverHook': 'tf.estimator.CheckpointSaverHook',
+ 'tf.train.CheckpointSaverListener': 'tf.estimator.CheckpointSaverListener',
'tf.train.ChiefSessionCreator': 'tf.compat.v1.train.ChiefSessionCreator',
+ 'tf.train.FeedFnHook': 'tf.estimator.FeedFnHook',
+ 'tf.train.FinalOpsHook': 'tf.estimator.FinalOpsHook',
'tf.train.FtrlOptimizer': 'tf.compat.v1.train.FtrlOptimizer',
+ 'tf.train.GlobalStepWaiterHook': 'tf.estimator.GlobalStepWaiterHook',
'tf.train.GradientDescentOptimizer': 'tf.compat.v1.train.GradientDescentOptimizer',
+ 'tf.train.LoggingTensorHook': 'tf.estimator.LoggingTensorHook',
'tf.train.LooperThread': 'tf.compat.v1.train.LooperThread',
'tf.train.MomentumOptimizer': 'tf.compat.v1.train.MomentumOptimizer',
'tf.train.MonitoredSession': 'tf.compat.v1.train.MonitoredSession',
'tf.train.MonitoredTrainingSession': 'tf.compat.v1.train.MonitoredTrainingSession',
- 'tf.train.NanLossDuringTrainingError': 'tf.compat.v1.train.NanLossDuringTrainingError',
+ 'tf.train.NanLossDuringTrainingError': 'tf.estimator.NanLossDuringTrainingError',
+ 'tf.train.NanTensorHook': 'tf.estimator.NanTensorHook',
'tf.train.NewCheckpointReader': 'tf.compat.v1.train.NewCheckpointReader',
'tf.train.Optimizer': 'tf.compat.v1.train.Optimizer',
- 'tf.train.ProfilerHook': 'tf.compat.v1.train.ProfilerHook',
+ 'tf.train.ProfilerHook': 'tf.estimator.ProfilerHook',
'tf.train.ProximalAdagradOptimizer': 'tf.compat.v1.train.ProximalAdagradOptimizer',
'tf.train.QueueRunner': 'tf.compat.v1.train.QueueRunner',
'tf.train.RMSPropOptimizer': 'tf.compat.v1.train.RMSPropOptimizer',
'tf.train.Saver': 'tf.compat.v1.train.Saver',
'tf.train.SaverDef': 'tf.compat.v1.train.SaverDef',
'tf.train.Scaffold': 'tf.compat.v1.train.Scaffold',
- 'tf.train.SecondOrStepTimer': 'tf.compat.v1.train.SecondOrStepTimer',
+ 'tf.train.SecondOrStepTimer': 'tf.estimator.SecondOrStepTimer',
+ 'tf.train.Server': 'tf.distribute.Server',
'tf.train.SessionCreator': 'tf.compat.v1.train.SessionCreator',
'tf.train.SessionManager': 'tf.compat.v1.train.SessionManager',
- 'tf.train.SessionRunArgs': 'tf.compat.v1.train.SessionRunArgs',
- 'tf.train.SessionRunContext': 'tf.compat.v1.train.SessionRunContext',
- 'tf.train.SessionRunValues': 'tf.compat.v1.train.SessionRunValues',
+ 'tf.train.SessionRunArgs': 'tf.estimator.SessionRunArgs',
+ 'tf.train.SessionRunContext': 'tf.estimator.SessionRunContext',
+ 'tf.train.SessionRunHook': 'tf.estimator.SessionRunHook',
+ 'tf.train.SessionRunValues': 'tf.estimator.SessionRunValues',
'tf.train.SingularMonitoredSession': 'tf.compat.v1.train.SingularMonitoredSession',
+ 'tf.train.StepCounterHook': 'tf.estimator.StepCounterHook',
+ 'tf.train.StopAtStepHook': 'tf.estimator.StopAtStepHook',
+ 'tf.train.SummarySaverHook': 'tf.estimator.SummarySaverHook',
'tf.train.Supervisor': 'tf.compat.v1.train.Supervisor',
'tf.train.SyncReplicasOptimizer': 'tf.compat.v1.train.SyncReplicasOptimizer',
'tf.train.VocabInfo': 'tf.estimator.VocabInfo',
diff --git a/tensorflow/tools/compatibility/reorders_v2.py b/tensorflow/tools/compatibility/reorders_v2.py
new file mode 100644
index 0000000..1c9fb92
--- /dev/null
+++ b/tensorflow/tools/compatibility/reorders_v2.py
@@ -0,0 +1,116 @@
+# 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.
+# ==============================================================================
+# pylint: disable=line-too-long
+"""List of renames to apply when converting from TF 1.0 to TF 2.0.
+
+THIS FILE IS AUTOGENERATED: To update, please run:
+ bazel build tensorflow/tools/compatibility/update:generate_v2_reorders_map
+ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
+This file should be updated whenever a function is added to
+self.reordered_function_names in tf_upgrade_v2.py.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+reorders = {
+ 'tf.argmax': ['input', 'axis', 'name', 'dimension', 'output_type'],
+ 'tf.argmin': ['input', 'axis', 'name', 'dimension', 'output_type'],
+ 'tf.batch_to_space': ['input', 'crops', 'block_size', 'name'],
+ 'tf.boolean_mask': ['tensor', 'mask', 'name', 'axis'],
+ 'tf.confusion_matrix': ['labels', 'predictions', 'num_classes', 'dtype', 'name', 'weights'],
+ 'tf.convert_to_tensor': ['value', 'dtype', 'name', 'preferred_dtype'],
+ 'tf.decode_csv': ['records', 'record_defaults', 'field_delim', 'use_quote_delim', 'name', 'na_value', 'select_cols'],
+ 'tf.depth_to_space': ['input', 'block_size', 'name', 'data_format'],
+ 'tf.feature_column.categorical_column_with_vocabulary_file': ['key', 'vocabulary_file', 'vocabulary_size', 'num_oov_buckets', 'default_value', 'dtype'],
+ 'tf.io.decode_csv': ['records', 'record_defaults', 'field_delim', 'use_quote_delim', 'name', 'na_value', 'select_cols'],
+ 'tf.io.parse_example': ['serialized', 'features', 'name', 'example_names'],
+ 'tf.io.parse_single_example': ['serialized', 'features', 'name', 'example_names'],
+ 'tf.io.serialize_many_sparse': ['sp_input', 'name', 'out_type'],
+ 'tf.io.serialize_sparse': ['sp_input', 'name', 'out_type'],
+ 'tf.linalg.norm': ['tensor', 'ord', 'axis', 'keepdims', 'name', 'keep_dims'],
+ 'tf.math.argmax': ['input', 'axis', 'name', 'dimension', 'output_type'],
+ 'tf.math.argmin': ['input', 'axis', 'name', 'dimension', 'output_type'],
+ 'tf.math.confusion_matrix': ['labels', 'predictions', 'num_classes', 'dtype', 'name', 'weights'],
+ 'tf.math.in_top_k': ['predictions', 'targets', 'k', 'name'],
+ 'tf.math.reduce_all': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_any': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_logsumexp': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_max': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_mean': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_min': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_prod': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.math.reduce_sum': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.multinomial': ['logits', 'num_samples', 'seed', 'name', 'output_dtype'],
+ 'tf.nn.convolution': ['input', 'filter', 'padding', 'strides', 'dilation_rate', 'name', 'data_format'],
+ 'tf.nn.crelu': ['features', 'name', 'axis'],
+ 'tf.nn.depth_to_space': ['input', 'block_size', 'name', 'data_format'],
+ 'tf.nn.depthwise_conv2d': ['input', 'filter', 'strides', 'padding', 'rate', 'name', 'data_format'],
+ 'tf.nn.embedding_lookup': ['params', 'ids', 'partition_strategy', 'name', 'validate_indices', 'max_norm'],
+ 'tf.nn.embedding_lookup_sparse': ['params', 'sp_ids', 'sp_weights', 'partition_strategy', 'name', 'combiner', 'max_norm'],
+ 'tf.nn.in_top_k': ['predictions', 'targets', 'k', 'name'],
+ 'tf.nn.moments': ['x', 'axes', 'shift', 'name', 'keep_dims'],
+ 'tf.nn.pool': ['input', 'window_shape', 'pooling_type', 'padding', 'dilation_rate', 'strides', 'name', 'data_format'],
+ 'tf.nn.separable_conv2d': ['input', 'depthwise_filter', 'pointwise_filter', 'strides', 'padding', 'rate', 'name', 'data_format'],
+ 'tf.nn.space_to_batch': ['input', 'paddings', 'block_size', 'name'],
+ 'tf.nn.space_to_depth': ['input', 'block_size', 'name', 'data_format'],
+ 'tf.nn.weighted_moments': ['x', 'axes', 'frequency_weights', 'name', 'keep_dims'],
+ 'tf.norm': ['tensor', 'ord', 'axis', 'keepdims', 'name', 'keep_dims'],
+ 'tf.pad': ['tensor', 'paddings', 'mode', 'name', 'constant_values'],
+ 'tf.parse_example': ['serialized', 'features', 'name', 'example_names'],
+ 'tf.parse_single_example': ['serialized', 'features', 'name', 'example_names'],
+ 'tf.quantize_v2': ['input', 'min_range', 'max_range', 'T', 'mode', 'name', 'round_mode'],
+ 'tf.random.multinomial': ['logits', 'num_samples', 'seed', 'name', 'output_dtype'],
+ 'tf.random.poisson': ['lam', 'shape', 'dtype', 'seed', 'name'],
+ 'tf.random_poisson': ['lam', 'shape', 'dtype', 'seed', 'name'],
+ 'tf.reduce_all': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_any': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_join': ['inputs', 'axis', 'keep_dims', 'separator', 'name', 'reduction_indices'],
+ 'tf.reduce_logsumexp': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_max': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_mean': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_min': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_prod': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reduce_sum': ['input_tensor', 'axis', 'keepdims', 'name', 'reduction_indices', 'keep_dims'],
+ 'tf.reverse_sequence': ['input', 'seq_lengths', 'seq_axis', 'batch_axis', 'name', 'seq_dim', 'batch_dim'],
+ 'tf.serialize_many_sparse': ['sp_input', 'name', 'out_type'],
+ 'tf.serialize_sparse': ['sp_input', 'name', 'out_type'],
+ 'tf.shape': ['input', 'name', 'out_type'],
+ 'tf.size': ['input', 'name', 'out_type'],
+ 'tf.space_to_batch': ['input', 'paddings', 'block_size', 'name'],
+ 'tf.space_to_depth': ['input', 'block_size', 'name', 'data_format'],
+ 'tf.sparse.add': ['a', 'b', 'threshold', 'thresh'],
+ 'tf.sparse.concat': ['axis', 'sp_inputs', 'name', 'expand_nonconcat_dim', 'concat_dim'],
+ 'tf.sparse.reduce_max': ['sp_input', 'axis', 'keepdims', 'reduction_axes', 'keep_dims'],
+ 'tf.sparse.segment_mean': ['data', 'indices', 'segment_ids', 'name', 'num_segments'],
+ 'tf.sparse.segment_sqrt_n': ['data', 'indices', 'segment_ids', 'name', 'num_segments'],
+ 'tf.sparse.segment_sum': ['data', 'indices', 'segment_ids', 'name', 'num_segments'],
+ 'tf.sparse.split': ['keyword_required', 'sp_input', 'num_split', 'axis', 'name', 'split_dim'],
+ 'tf.sparse_add': ['a', 'b', 'threshold', 'thresh'],
+ 'tf.sparse_concat': ['axis', 'sp_inputs', 'name', 'expand_nonconcat_dim', 'concat_dim'],
+ 'tf.sparse_matmul': ['a', 'b', 'transpose_a', 'transpose_b', 'a_is_sparse', 'b_is_sparse', 'name'],
+ 'tf.sparse_reduce_max': ['sp_input', 'axis', 'keepdims', 'reduction_axes', 'keep_dims'],
+ 'tf.sparse_segment_mean': ['data', 'indices', 'segment_ids', 'name', 'num_segments'],
+ 'tf.sparse_segment_sqrt_n': ['data', 'indices', 'segment_ids', 'name', 'num_segments'],
+ 'tf.sparse_segment_sum': ['data', 'indices', 'segment_ids', 'name', 'num_segments'],
+ 'tf.sparse_split': ['keyword_required', 'sp_input', 'num_split', 'axis', 'name', 'split_dim'],
+ 'tf.strings.length': ['input', 'name', 'unit'],
+ 'tf.strings.reduce_join': ['inputs', 'axis', 'keep_dims', 'separator', 'name', 'reduction_indices'],
+ 'tf.strings.substr': ['input', 'pos', 'len', 'name', 'unit'],
+ 'tf.substr': ['input', 'pos', 'len', 'name', 'unit'],
+ 'tf.transpose': ['a', 'perm', 'name', 'conjugate'],
+ 'tf.tuple': ['tensors', 'name', 'control_inputs'],
+ 'tf.while_loop': ['cond', 'body', 'loop_vars', 'shape_invariants', 'parallel_iterations', 'back_prop', 'swap_memory', 'name', 'maximum_iterations', 'return_same_structure']
+}
diff --git a/tensorflow/tools/compatibility/testdata/test_file_v0_11.py b/tensorflow/tools/compatibility/testdata/test_file_v0_11.py
index 68ba7a2..917236d 100644
--- a/tensorflow/tools/compatibility/testdata/test_file_v0_11.py
+++ b/tensorflow/tools/compatibility/testdata/test_file_v0_11.py
@@ -34,6 +34,7 @@
a unit test if the converter is successful.
"""
+ @test_util.run_v1_only("b/120545219")
def testArgRenames(self):
with self.cached_session():
@@ -97,6 +98,7 @@
tf.expand_dims([[1, 2], [3, 4]], axis=1).eval(),
[[[1, 2]], [[3, 4]]])
+ @test_util.run_v1_only("b/120545219")
def testArgMinMax(self):
with self.cached_session():
self.assertAllEqual(
@@ -112,6 +114,7 @@
tf.argmax([[1, 2, 3], [4, 1, 0]], dimension=0).eval(),
[1, 0, 0])
+ @test_util.run_v1_only("b/120545219")
def testExpandAndSqueeze(self):
with self.cached_session():
@@ -139,6 +142,7 @@
[[1, 2, 3]], dim=1), squeeze_dims=[1]).eval(),
a)
+ @test_util.run_v1_only("b/120545219")
def testArithmeticRenames(self):
with self.cached_session() as s:
stuff = tf.split(1, 2, [[1, 2, 3, 4], [4, 5, 6, 7]])
@@ -163,6 +167,7 @@
# # TODO(aselle): (tf.batch_*)
# ]
+ @test_util.run_v1_only("b/120545219")
def testBatchAndSvd(self):
with self.cached_session():
mat = [[1., 2.], [2., 3.]]
@@ -174,6 +179,7 @@
tf.svd(mat, False, True).eval(),
tf.svd(mat, compute_uv=False, full_matrices=True).eval())
+ @test_util.run_v1_only("b/120545219")
def testCrossEntropy(self):
# TODO(aselle): Test sparse_softmax_...
with self.cached_session():
@@ -190,6 +196,7 @@
tf.nn.sigmoid_cross_entropy_with_logits(
labels=labels, logits=logits).eval())
+ @test_util.run_v1_only("b/120545219")
def testVariables(self):
with self.cached_session() as s:
@@ -200,6 +207,7 @@
_ = [v.name for v in tf.all_variables()]
_ = [v.name for v in tf.local_variables()]
+ @test_util.run_v1_only("b/120545219")
def testSummaries(self):
with self.cached_session() as s:
var = tf.Variable([1, 2, 3], dtype=tf.float32)
diff --git a/tensorflow/tools/compatibility/testdata/test_file_v1_12.py b/tensorflow/tools/compatibility/testdata/test_file_v1_12.py
index fd68878..5ce4dd4 100644
--- a/tensorflow/tools/compatibility/testdata/test_file_v1_12.py
+++ b/tensorflow/tools/compatibility/testdata/test_file_v1_12.py
@@ -28,11 +28,13 @@
def setUp(self):
tf.enable_eager_execution()
+ @test_util.run_v1_only("b/120545219")
def testRenames(self):
with self.cached_session():
self.assertAllClose(1.04719755, tf.acos(0.5))
self.assertAllClose(0.5, tf.rsqrt(4.0))
+ @test_util.run_v1_only("b/120545219")
def testSerializeSparseTensor(self):
sp_input = tf.SparseTensor(
indices=tf.constant([[1]], dtype=tf.int64),
@@ -44,6 +46,7 @@
self.assertEqual((3,), serialized_sp.shape)
self.assertTrue(serialized_sp[0].numpy()) # check non-empty
+ @test_util.run_v1_only("b/120545219")
def testSerializeManySparse(self):
sp_input = tf.SparseTensor(
indices=tf.constant([[0, 1]], dtype=tf.int64),
@@ -55,6 +58,7 @@
sp_input, 'serialize_name', tf.string)
self.assertEqual((1, 3), serialized_sp.shape)
+ @test_util.run_v1_only("b/120545219")
def testArgMaxMin(self):
self.assertAllClose(
[1],
diff --git a/tensorflow/tools/compatibility/tf_upgrade_v2.py b/tensorflow/tools/compatibility/tf_upgrade_v2.py
index d19eb13..427e22b 100644
--- a/tensorflow/tools/compatibility/tf_upgrade_v2.py
+++ b/tensorflow/tools/compatibility/tf_upgrade_v2.py
@@ -20,6 +20,7 @@
from tensorflow.tools.compatibility import ast_edits
from tensorflow.tools.compatibility import renames_v2
+from tensorflow.tools.compatibility import reorders_v2
class TFAPIChangeSpec(ast_edits.APIChangeSpec):
@@ -35,6 +36,18 @@
"tf.argmax": {
"dimension": "axis",
},
+ "tf.arg_min": {
+ "dimension": "axis",
+ },
+ "tf.arg_max": {
+ "dimension": "axis",
+ },
+ "tf.math.argmin": {
+ "dimension": "axis",
+ },
+ "tf.math.argmax": {
+ "dimension": "axis",
+ },
"tf.image.crop_and_resize": {
"box_ind": "box_indices",
},
@@ -50,6 +63,12 @@
"tf.batch_to_space": {
"block_size": "block_shape",
},
+ "tf.space_to_batch": {
+ "block_size": "block_shape",
+ },
+ "tf.nn.space_to_batch": {
+ "block_size": "block_shape",
+ },
"tf.constant": {
"verify_shape": "verify_shape_is_now_always_true",
},
@@ -62,6 +81,15 @@
"tf.linalg.l2_normalize": {
"dim": "axis",
},
+ "tf.linalg.norm": {
+ "keep_dims": "keepdims",
+ },
+ "tf.norm": {
+ "keep_dims": "keepdims",
+ },
+ "tf.load_file_system_library": {
+ "library_filename": "library_location",
+ },
"tf.math.count_nonzero": {
"input_tensor": "input",
"keep_dims": "keepdims",
@@ -95,6 +123,9 @@
"tf.nn.separable_conv2d": {
"rate": "dilations"
},
+ "tf.nn.depthwise_conv2d": {
+ "rate": "dilations"
+ },
"tf.nn.softmax": {
"dim": "axis"
},
@@ -113,14 +144,35 @@
},
"tf.sparse.concat": {
"concat_dim": "axis",
+ "expand_nonconcat_dim": "expand_nonconcat_dims",
},
"tf.sparse_concat": {
"concat_dim": "axis",
+ "expand_nonconcat_dim": "expand_nonconcat_dims",
},
"tf.sparse.split": {
"split_dim": "axis",
},
- "tf.max_pool_with_argmax": {
+ "tf.sparse_split": {
+ "split_dim": "axis",
+ },
+ "tf.sparse.reduce_max": {
+ "reduction_axes": "axis",
+ "keep_dims": "keepdims",
+ },
+ "tf.sparse_reduce_max": {
+ "reduction_axes": "axis",
+ "keep_dims": "keepdims",
+ },
+ "tf.sparse.reduce_sum": {
+ "reduction_axes": "axis",
+ "keep_dims": "keepdims",
+ },
+ "tf.sparse_reduce_sum": {
+ "reduction_axes": "axis",
+ "keep_dims": "keepdims",
+ },
+ "tf.nn.max_pool_with_argmax": {
"Targmax": "output_dtype",
},
"tf.multinomial": {
@@ -129,6 +181,10 @@
"tf.random.multinomial": {
"output_dtype": "dtype",
},
+ "tf.reverse_sequence": {
+ "seq_dim": "seq_axis",
+ "batch_dim": "batch_axis",
+ },
"tf.nn.batch_norm_with_global_normalization": {
"t": "input",
"m": "mean",
@@ -147,6 +203,10 @@
"tf.ones_like": {
"tensor": "input",
},
+ "tf.nn.conv2d_transpose": {
+ "value": "input",
+ "filter": "filters",
+ },
"tf.nn.conv3d_transpose": {
"value": "input",
"filter": "filters",
@@ -283,6 +343,9 @@
"tf.squeeze": {
"squeeze_dims": "axis",
},
+ "tf.nn.weighted_moments": {
+ "keep_dims": "keepdims"
+ },
}
# pylint: disable=line-too-long
@@ -293,6 +356,10 @@
self.manual_symbol_renames = {
"tf.batch_to_space_nd":
"tf.batch_to_space",
+ "tf.space_to_batch_nd":
+ "tf.space_to_batch",
+ "tf.nn.space_to_batch":
+ "tf.space_to_batch",
"tf.extract_image_patches":
"tf.image.extract_image_patches",
"tf.gfile.Copy":
@@ -419,8 +486,14 @@
"tf.sparse.split",
"tf.sparse_matmul":
"tf.linalg.matmul",
+ "tf.sparse_reduce_sum":
+ "tf.sparse.reduce_sum",
+ "tf.sparse_reduce_max":
+ "tf.sparse.reduce_max",
"tf.random.stateless_multinomial":
"tf.random.stateless_categorical",
+ "tf.substr":
+ "tf.strings.substr",
"tf.string_to_hash_bucket":
"tf.strings.to_hash_bucket",
"tf.string_to_number":
@@ -473,6 +546,10 @@
"tf.argmax",
"tf.arg_min":
"tf.argmin",
+ # tf.nn.ctc_loss is still available in 2.0 but behavior
+ # changed significantly.
+ "tf.nn.ctc_loss":
+ "tf.compat.v1.nn.ctc_loss",
}
# pylint: enable=line-too-long
@@ -483,187 +560,94 @@
# Variables that should be changed to functions.
self.change_to_function = {}
+ # pylint: disable=line-too-long
+ # This list should just contain names of functions that had
+ # their arguments reordered. After adding a function name to the list
+ # run the following to update reorders_v2.py:
+ # bazel build tensorflow/tools/compatibility/update:generate_v2_reorders_map
+ # bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
+ # pylint: enable=line-too-long
+ self.reordered_function_names = {
+ "tf.io.serialize_sparse",
+ "tf.io.serialize_many_sparse",
+ "tf.argmax",
+ "tf.argmin",
+ "tf.batch_to_space",
+ "tf.nn.space_to_batch",
+ "tf.boolean_mask",
+ "tf.convert_to_tensor",
+ "tf.nn.moments",
+ "tf.nn.convolution",
+ "tf.nn.crelu",
+ "tf.nn.weighted_moments",
+ "tf.nn.pool",
+ "tf.nn.separable_conv2d",
+ "tf.nn.depthwise_conv2d",
+ "tf.multinomial",
+ "tf.random.multinomial",
+ "tf.pad",
+ "tf.quantize_v2",
+ "tf.feature_column.categorical_column_with_vocabulary_file",
+ "tf.shape",
+ "tf.size",
+ "tf.random.poisson",
+ "tf.sparse.add",
+ "tf.sparse_add",
+ "tf.sparse.concat",
+ "tf.sparse_concat",
+ "tf.sparse.segment_mean",
+ "tf.sparse.segment_sqrt_n",
+ "tf.sparse.segment_sum",
+ "tf.sparse_matmul",
+ "tf.sparse.reduce_max",
+ "tf.sparse_reduce_max",
+ "tf.io.decode_csv",
+ "tf.strings.length",
+ "tf.strings.reduce_join",
+ "tf.strings.substr",
+ "tf.substr",
+ "tf.transpose",
+ "tf.tuple",
+ "tf.parse_example",
+ "tf.parse_single_example",
+ "tf.io.parse_example",
+ "tf.io.parse_single_example",
+ "tf.while_loop",
+ "tf.reduce_all",
+ "tf.math.reduce_all",
+ "tf.reduce_any",
+ "tf.math.reduce_any",
+ "tf.reduce_min",
+ "tf.math.reduce_min",
+ "tf.reduce_max",
+ "tf.math.reduce_max",
+ "tf.reduce_sum",
+ "tf.math.reduce_sum",
+ "tf.reduce_mean",
+ "tf.math.reduce_mean",
+ "tf.reduce_prod",
+ "tf.math.reduce_prod",
+ "tf.reduce_logsumexp",
+ "tf.math.reduce_logsumexp",
+ "tf.reduce_join",
+ "tf.confusion_matrix",
+ "tf.math.confusion_matrix",
+ "tf.math.in_top_k",
+ "tf.nn.depth_to_space",
+ "tf.nn.embedding_lookup",
+ "tf.nn.embedding_lookup_sparse",
+ "tf.nn.in_top_k",
+ "tf.nn.space_to_depth",
+ "tf.linalg.norm",
+ "tf.norm",
+ "tf.reverse_sequence",
+ "tf.sparse_split",
+ }
+
# Functions that were reordered should be changed to the new keyword args
# for safety, if positional arguments are used. If you have reversed the
# positional arguments yourself, this could do the wrong thing.
- # IMPORTANT: order here should correspond to OLD argument order.
- # We just prepend "arg_name=" to all arguments in function calls.
- self.function_reorders = {
- "tf.io.serialize_sparse": ["sp_input", "name", "out_type"],
- "tf.io.serialize_many_sparse": ["sp_input", "name", "out_type"],
- "tf.argmax": ["input", "axis", "name", "axis", "output_type"],
- "tf.argmin": ["input", "axis", "name", "axis", "output_type"],
- "tf.batch_to_space": ["input", "crops", "block_size", "name"],
- "tf.boolean_mask": ["tensor", "mask", "name", "axis"],
- "tf.convert_to_tensor": ["value", "dtype", "name", "preferred_dtype"],
- "tf.nn.moments": ["x", "axes", "shift", "keepdims", "name"],
- "tf.nn.convolution": [
- "input", "filter", "padding", "strides", "dilation_rate", "name",
- "data_format"
- ],
- "tf.nn.crelu": ["features", "name", "axis"],
- "tf.nn.weighted_moments": [
- "x", "axes", "frequency_weights", "name", "keep_dims"
- ],
- "tf.nn.pool": [
- "input", "window_shape", "pooling_type", "padding", "dilation_rate",
- "strides", "name", "data_format"
- ],
- "tf.nn.separable_conv2d": [
- "input", "depthwise_filter", "pointwise_filter", "strides",
- "padding", "rate", "name", "data_format"
- ],
- "tf.nn.depthwise_conv2d": [
- "input", "filter", "strides", "padding", "rate", "name",
- "data_format"
- ],
- "tf.multinomial": [
- "logits", "num_samples", "seed", "name", "output_dtype"
- ],
- "tf.random.multinomial": [
- "logits", "num_samples", "seed", "name", "output_dtype"
- ],
- "tf.pad": ["tensor", "paddings", "mode", "name", "constant_values"],
- "tf.quantize_v2": [
- "input", "min_range", "max_range", "T", "mode", "name", "round_mode"
- ],
- "tf.feature_column.categorical_column_with_vocabulary_file": [
- "key", "vocabulary_file", "vocabulary_size", "num_oov_buckets",
- "default_value", "dtype"
- ],
- "tf.shape": ["input", "name", "out_type"],
- "tf.size": ["input", "name", "out_type"],
- "tf.random.poisson": ["lam", "shape", "dtype", "seed", "name"],
- "tf.sparse.add": ["a", "b", "thresh"],
- "tf.sparse_add": ["a", "b", "thresh"],
- "tf.sparse.concat": [
- "axis", "sp_inputs", "name", "expand_nonconcat_dim", "concat_dim"
- ],
- "tf.sparse_concat": [
- "axis", "sp_inputs", "name", "expand_nonconcat_dim", "concat_dim"
- ],
- "tf.sparse.segment_mean": [
- "data", "indices", "segment_ids", "name", "num_segments"
- ],
- "tf.sparse.segment_sqrt_n": [
- "data", "indices", "segment_ids", "name", "num_segments"
- ],
- "tf.sparse.segment_sum": [
- "data", "indices", "segment_ids", "name", "num_segments"
- ],
- "tf.sparse_matmul": [
- "a", "b", "transpose_a", "transpose_b", "a_is_sparse",
- "b_is_sparse", "name"
- ],
- "tf.io.decode_csv": [
- "records",
- "record_defaults",
- "field_delim",
- "use_quote_delim",
- "name",
- "na_value",
- "select_cols",
- ],
- "tf.strings.substr": ["input", "pos", "len", "name", "unit"],
- "tf.strings.reduce_join": [
- "input", "axis", "keep_dims", "separator", "name",
- "reduction_indices"
- ],
- "tf.strings.length": ["input", "name", "unit"],
- "tf.transpose": ["a", "perm", "name", "conjugate"],
- "tf.tuple": ["tensors", "name", "control_inputs"],
- "tf.parse_example": [
- "serialized", "features", "name", "example_names"
- ],
- "tf.parse_single_example": [
- "serialized", "features", "name", "example_names"
- ],
- "tf.io.parse_example": [
- "serialized", "features", "name", "example_names"
- ],
- "tf.io.parse_single_example": [
- "serialized", "features", "name", "example_names"
- ],
- "tf.while_loop": [
- "cond", "body", "loop_vars", "shape_invariants",
- "parallel_iterations", "back_prop", "swap_memory", "name",
- "maximum_iterations", "return_same_structure"
- ],
- "tf.reduce_all": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_all": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_any": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_any": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_min": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_min": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_max": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_max": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_sum": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_sum": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_mean": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_mean": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_prod": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_prod": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_logsumexp": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.math.reduce_logsumexp": [
- "input_tensor", "axis", "keepdims", "name", "reduction_indices",
- "keep_dims"
- ],
- "tf.reduce_join": [
- "input", "axis", "keep_dims", "separator", "name",
- "reduction_indices"
- ],
- "tf.confusion_matrix": [
- "labels", "predictions", "num_classes", "dtype", "name", "weights"
- ],
- "tf.math.confusion_matrix": [
- "labels", "predictions", "num_classes", "dtype", "name", "weights"
- ]
- }
+ self.function_reorders = reorders_v2.reorders
# Specially handled functions.
self.function_handle = {
@@ -711,30 +695,60 @@
" they may already have been correct)."
)
+ deprecate_partition_strategy_comment = (
+ "WARNING: `partition_strategy` has been removed from `%s` "
+ " The 'div' strategy is used by default.")
+
# Function warnings. <function name> placeholder inside warnings will be
# replaced by function name.
self.function_warnings = {
- "tf.assert_greater": assert_return_type_comment,
- "tf.assert_equal": assert_return_type_comment,
- "tf.assert_less": assert_return_type_comment,
- "tf.assert_rank": assert_rank_comment,
- "tf.debugging.assert_equal": assert_return_type_comment,
- "tf.debugging.assert_greater": assert_return_type_comment,
- "tf.debugging.assert_greater_equal": assert_return_type_comment,
- "tf.debugging.assert_integer": assert_return_type_comment,
- "tf.debugging.assert_less": assert_return_type_comment,
- "tf.debugging.assert_less_equal": assert_return_type_comment,
- "tf.debugging.assert_near": assert_return_type_comment,
- "tf.debugging.assert_negative": assert_return_type_comment,
- "tf.debugging.assert_non_negative": assert_return_type_comment,
- "tf.debugging.assert_non_positive": assert_return_type_comment,
- "tf.debugging.assert_none_equal": assert_return_type_comment,
- "tf.debugging.assert_positive": assert_return_type_comment,
- "tf.debugging.assert_rank": assert_rank_comment,
- "tf.debugging.assert_rank_at_least": assert_rank_comment,
- "tf.debugging.assert_rank_in": assert_rank_comment,
- "tf.flags": "tf.flags has been removed, please use the argparse or absl"
- " module if you need command line parsing.",
+ "tf.assert_greater":
+ assert_return_type_comment,
+ "tf.assert_equal":
+ assert_return_type_comment,
+ "tf.assert_less":
+ assert_return_type_comment,
+ "tf.assert_rank":
+ assert_rank_comment,
+ "tf.cond": "tf.cond no longer takes 'strict'. "
+ "Now 'strict' defaults to True."
+ "fn1/fn2 arguments are replaced by true_fn/false_fn.",
+ "tf.debugging.assert_equal":
+ assert_return_type_comment,
+ "tf.debugging.assert_greater":
+ assert_return_type_comment,
+ "tf.debugging.assert_greater_equal":
+ assert_return_type_comment,
+ "tf.debugging.assert_integer":
+ assert_return_type_comment,
+ "tf.debugging.assert_less":
+ assert_return_type_comment,
+ "tf.debugging.assert_less_equal":
+ assert_return_type_comment,
+ "tf.debugging.assert_near":
+ assert_return_type_comment,
+ "tf.debugging.assert_negative":
+ assert_return_type_comment,
+ "tf.debugging.assert_non_negative":
+ assert_return_type_comment,
+ "tf.debugging.assert_non_positive":
+ assert_return_type_comment,
+ "tf.debugging.assert_none_equal":
+ assert_return_type_comment,
+ "tf.debugging.assert_positive":
+ assert_return_type_comment,
+ "tf.debugging.assert_rank":
+ assert_rank_comment,
+ "tf.debugging.assert_rank_at_least":
+ assert_rank_comment,
+ "tf.debugging.assert_rank_in":
+ assert_rank_comment,
+ "tf.device": "tf.device no longer takes function as an argument. "
+ "'devide_name_or_function' argument has been renamed to "
+ "'device_name'.",
+ "tf.flags":
+ "tf.flags has been removed, please use the argparse or absl"
+ " module if you need command line parsing.",
"tf.train.exponential_decay":
decay_function_comment,
"tf.train.piecewise_constant_decay":
@@ -769,24 +783,63 @@
default_loss_reduction_changed,
"tf.estimator.BaselineRegressor":
default_loss_reduction_changed,
+ "tf.hessians": "tf.hessians no longer takes "
+ "'colocate_gradients_with_ops' argument. Also, "
+ "arguments have been reordered so that 'name' is the "
+ "last argument.",
"tf.nn.conv1d":
- "WARNING: use_cudnn_on_gpu argument has been removed and \"value\" was "
- "renamed to \"input\"",
+ "WARNING: use_cudnn_on_gpu argument has been removed and \"value\""
+ " was renamed to \"input\"",
"tf.nn.conv2d":
- "WARNING: use_cudnn_on_gpu argument has been removed and \"filter\" "
- "was renamed to \"filters\"",
+ "WARNING: use_cudnn_on_gpu argument has been removed and "
+ "\"filter\" was renamed to \"filters\"",
"tf.nn.conv2d_backprop_filter":
- "WARNING: use_cudnn_on_gpu argument has been removed",
+ "WARNING: use_cudnn_on_gpu argument has been removed",
"tf.nn.conv2d_backprop_input":
- "WARNING: use_cudnn_on_gpu argument has been removed and \"filter\" "
- "was renamed to \"filters\"",
+ "WARNING: use_cudnn_on_gpu argument has been removed and "
+ "\"filter\" was renamed to \"filters\"",
"tf.nn.erosion2d":
- "WARNING: <function name> now requires a data_format argument",
+ "WARNING: <function name> now requires a data_format argument",
"tf.nn.nce_loss":
- "WARNING: `partition_strategy` has been removed from `tf.nn.nce_loss` "
- " The 'div' strategy is used by default.",
- "tf.zeros_like": tf_01s_like_no_optimize_comment,
- "tf.ones_like": tf_01s_like_no_optimize_comment,
+ deprecate_partition_strategy_comment % "tf.nn.nce_loss",
+ "tf.nn.safe_embedding_lookup_sparse":
+ deprecate_partition_strategy_comment %
+ "tf.nn.safe_embedding_lookup_sparse",
+ "tf.nn.sampled_softmax_loss":
+ deprecate_partition_strategy_comment % "tf.nn.sampled_softmax_loss",
+ "tf.zeros_like":
+ tf_01s_like_no_optimize_comment,
+ "tf.ones_like":
+ tf_01s_like_no_optimize_comment,
+ "tf.nn.embedding_lookup":
+ "WARNING: validate_indices argument has been removed.",
+ "tf.while_loop":
+ "tf.while_loop no longer takes 'return_same_structure' argument. "
+ "'return_same_structure' now defaults to True. Also, 'name'"
+ "argument is now the last argument.",
+ "tf.image.sample_distorted_bounding_box":
+ "tf.image.sample_distorted_bounding_box no longer takes 'seed2' "
+ "argument.",
+ "tf.nn.ctc_beam_search_decoder":
+ "tf.nn.ctc_beam_search_decoder no longer takes 'merge_repeated' "
+ "argument. 'merge_repeated' now defaults to False.",
+ "tf.nn.fractional_avg_pool":
+ "tf.nn.fractional_avg_pool no longer takes 'seed2' and "
+ "'deterministic' arguments. Now it takes a single 'seed' arg. If "
+ "'seed' is zero, the execution is random and deterministic "
+ "otherwise",
+ "tf.nn.fractional_max_pool":
+ "tf.nn.fractional_max_pool no longer takes 'seed2' and "
+ "'deterministic' arguments. Now it takes a single 'seed' arg. If "
+ "'seed' is zero, the execution is random and deterministic "
+ "otherwise",
+ "tf.nn.softmax_cross_entropy_with_logits":
+ "tf.nn.softmax_cross_entropy_with_logits behavior has changed. "
+ "'labels' needs to be wrapped with tf.stop_gradient to keep the "
+ "old behavior. Also, 'dim' argument has been renamed to 'axis'.",
+ "tf.test.assert_equal_graph_def":
+ "tf.assert_equal_graph_def no longer takes 'checkpoint_v2' "
+ "argument. 'checkpoint_v2' now defaults to True.",
}
self.symbol_renames = {
diff --git a/tensorflow/tools/compatibility/tf_upgrade_v2_test.py b/tensorflow/tools/compatibility/tf_upgrade_v2_test.py
index bc92cab..d5428e7 100644
--- a/tensorflow/tools/compatibility/tf_upgrade_v2_test.py
+++ b/tensorflow/tools/compatibility/tf_upgrade_v2_test.py
@@ -18,6 +18,7 @@
from __future__ import division
from __future__ import print_function
+import inspect
import os
import tempfile
@@ -25,7 +26,6 @@
import tensorflow as tf
# OSS TF V2 import placeholder.
-
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test as test_lib
from tensorflow.python.util import tf_decorator
@@ -37,30 +37,22 @@
from tensorflow.tools.compatibility import tf_upgrade_v2
-_TENSORFLOW_API_ATTR_V1 = (
- tf_export.API_ATTRS_V1[tf_export.TENSORFLOW_API_NAME].names)
-_TENSORFLOW_API_ATTR = tf_export.API_ATTRS[tf_export.TENSORFLOW_API_NAME].names
-_ESTIMATOR_API_ATTR_V1 = (
- tf_export.API_ATTRS_V1[tf_export.ESTIMATOR_API_NAME].names)
-_ESTIMATOR_API_ATTR = tf_export.API_ATTRS[tf_export.ESTIMATOR_API_NAME].names
+def get_symbol_for_name(root, name):
+ name_parts = name.split(".")
+ symbol = root
+ # Iterate starting with second item since 1st item is "tf.".
+ for part in name_parts[1:]:
+ symbol = getattr(symbol, part)
+ return symbol
-def get_v1_names(symbol):
- names_v1 = []
- if hasattr(symbol, _TENSORFLOW_API_ATTR_V1):
- names_v1.extend(getattr(symbol, _TENSORFLOW_API_ATTR_V1))
- if hasattr(symbol, _ESTIMATOR_API_ATTR_V1):
- names_v1.extend(getattr(symbol, _ESTIMATOR_API_ATTR_V1))
- return names_v1
-
-
-def get_v2_names(symbol):
- names_v2 = set()
- if hasattr(symbol, _TENSORFLOW_API_ATTR):
- names_v2.update(getattr(symbol, _TENSORFLOW_API_ATTR))
- if hasattr(symbol, _ESTIMATOR_API_ATTR):
- names_v2.update(getattr(symbol, _ESTIMATOR_API_ATTR))
- return list(names_v2)
+def get_args(symbol):
+ if hasattr(inspect, "signature"):
+ signature = inspect.signature(symbol)
+ # Ignore *args and **kwargs for now.
+ return [param.name for param in signature.parameters.values()
+ if param.kind == param.POSITIONAL_OR_KEYWORD]
+ return tf_inspect.getargspec(symbol)[0]
def get_func_and_args_from_str(call_str):
@@ -79,6 +71,7 @@
function_name = call_str[:call_str.find("(")]
args = call_str[open_paren_index+1:close_paren_index].split(",")
args = [arg.split("=")[0].strip() for arg in args]
+ args = [arg for arg in args if arg] # filter out empty strings
return function_name, args
@@ -99,7 +92,7 @@
def symbol_collector(unused_path, unused_parent, children):
for child in children:
_, attr = tf_decorator.unwrap(child[1])
- api_names_v2 = get_v2_names(attr)
+ api_names_v2 = tf_export.get_v2_names(attr)
for name in api_names_v2:
cls.v2_symbols["tf." + name] = attr
@@ -142,10 +135,12 @@
# Converts all symbols in the v1 namespace to the v2 namespace, raising
# an error if the target of the conversion is not in the v2 namespace.
+ # Please regenerate the renames file or edit any manual renames if this
+ # test fails.
def conversion_visitor(unused_path, unused_parent, children):
for child in children:
_, attr = tf_decorator.unwrap(child[1])
- api_names = get_v1_names(attr)
+ api_names = tf_export.get_v1_names(attr)
for name in api_names:
_, _, _, text = self._upgrade("tf." + name)
if (text and
@@ -160,20 +155,51 @@
visitor.private_map["tf.compat"] = ["v1", "v2"]
traverse.traverse(tf.compat.v1, visitor)
- def testKeywordArgNames(self):
- if not hasattr(tf.compat, "v2"):
- return
+ def testAllAPIV1(self):
+ collect = True
+ v1_symbols = set([])
- all_keyword_renames = (
- tf_upgrade_v2.TFAPIChangeSpec().function_keyword_renames)
- v2_name_exceptions = {"verify_shape_is_now_always_true"}
+ # Symbols which may be generated by the conversion script which do not exist
+ # in TF 1.x. This should be a very short list of symbols which are
+ # experimental in 1.x but stable for 2.x.
+ whitelisted_v2_only_symbols = set(["tf.saved_model.save"])
- # Visitor that verifies V1 argument names, converts to V2 and checks
- # V2 argument names.
+ # Converts all symbols in the v1 namespace to the v2 namespace, raising
+ # an error if the target of the conversion is not in the v1 namespace.
def conversion_visitor(unused_path, unused_parent, children):
for child in children:
_, attr = tf_decorator.unwrap(child[1])
- names_v1 = get_v1_names(attr)
+ api_names = tf_export.get_v1_names(attr)
+ for name in api_names:
+ if collect:
+ v1_symbols.add("tf." + name)
+ else:
+ _, _, _, text = self._upgrade("tf." + name)
+ if (text and
+ not text.startswith("tf.compat.v1") and
+ not text.startswith("tf.estimator") and
+ text not in v1_symbols and
+ text not in whitelisted_v2_only_symbols):
+ self.assertFalse(
+ True, "Symbol %s generated from %s not in v1 API" % (
+ text, name))
+
+ visitor = public_api.PublicAPIVisitor(conversion_visitor)
+ visitor.do_not_descend_map["tf"].append("contrib")
+ visitor.private_map["tf.compat"] = ["v1", "v2"]
+ traverse.traverse(tf.compat.v1, visitor)
+ collect = False
+ traverse.traverse(tf.compat.v1, visitor)
+
+ def testV1KeywordArgNames(self):
+ all_keyword_renames = (
+ tf_upgrade_v2.TFAPIChangeSpec().function_keyword_renames)
+
+ # Visitor that verifies V1 argument names.
+ def arg_test_visitor(unused_path, unused_parent, children):
+ for child in children:
+ _, attr = tf_decorator.unwrap(child[1])
+ names_v1 = tf_export.get_v1_names(attr)
for name in names_v1:
name = "tf.%s" % name
@@ -190,31 +216,118 @@
"%s not found in %s arguments: %s" %
(from_name, name, str(arg_names_v1)))
+ visitor = public_api.PublicAPIVisitor(arg_test_visitor)
+ visitor.do_not_descend_map["tf"].append("contrib")
+ visitor.private_map["tf.compat"] = ["v1", "v2"]
+ traverse.traverse(tf.compat.v1, visitor)
+
+ def testV2KeywordArgNames(self):
+ # This test converts a call of the form:
+ # tf.foo(arg1=0, arg2=1, ...)
+ # to 2.0. Then, checks that converted function has valid argument names.
+ if not hasattr(tf.compat, "v2"):
+ return
+ v2_arg_exceptions = {
+ "verify_shape_is_now_always_true",
+ # These arguments should not be used, they just specify
+ # that a function takes named arguments.
+ "keyword_required",
+ "_sentinel",
+ }
+ v1_name_exceptions = {
+ "tf.print", # requires print_function import
+ }
+ function_warnings = (
+ tf_upgrade_v2.TFAPIChangeSpec().function_warnings)
+ function_handles = (
+ tf_upgrade_v2.TFAPIChangeSpec().function_handle)
+ keyword_renames = (
+ tf_upgrade_v2.TFAPIChangeSpec().function_keyword_renames)
+
+ # Visitor that converts to V2 and checks V2 argument names.
+ def conversion_visitor(unused_path, unused_parent, children):
+ for child in children:
+ _, attr = tf_decorator.unwrap(child[1])
+ if not tf_inspect.isfunction(attr):
+ continue
+ names_v1 = tf_export.get_v1_names(attr)
+ arg_names_v1 = get_args(attr)
+
+ for name in names_v1:
+ tf_name = "tf.%s" % name
+ if tf_name in function_warnings or tf_name in function_handles:
+ continue # These require manual change
+ if tf_name in v1_name_exceptions:
+ continue
# Assert that arg names after converting to v2 are present in
# v2 function.
# 1. First, create an input of the form:
# tf.foo(arg1=val1, arg2=val2, ...)
args = ",".join(
["%s=%d" % (from_name, from_index)
- for from_index, from_name in enumerate(keyword_renames.keys())])
- text_input = "%s(%s)" % (name, args)
+ for from_index, from_name in enumerate(arg_names_v1)])
+ text_input = "%s(%s)" % (tf_name, args)
# 2. Convert the input to V2.
_, _, _, text = self._upgrade(text_input)
new_function_name, new_args = get_func_and_args_from_str(text)
+ if new_function_name == "tf.compat.v1.%s" % name:
+ if tf_name in keyword_renames:
+ # If we rename arguments, new function must be available in 2.0.
+ # We should not be using compat.v1 in this case.
+ self.assertFalse(
+ "Function '%s' is not in 2.0 when converting\n%s\nto\n%s" %
+ (new_function_name, text_input, text))
+ continue
# 3. Verify V2 function and arguments.
- # Note: If we rename arguments, new function must be available in 2.0.
- # We should not be using compat.v1 in this case.
- self.assertIn(new_function_name, self.v2_symbols)
- args_v2 = tf_inspect.getargspec(self.v2_symbols[new_function_name])[0]
- args_v2.extend(v2_name_exceptions)
+ args_v2 = get_args(self.v2_symbols[new_function_name])
+ args_v2.extend(v2_arg_exceptions)
for new_arg in new_args:
- self.assertIn(new_arg, args_v2)
+ self.assertIn(
+ new_arg, args_v2,
+ "Invalid argument '%s' in 2.0 when converting\n%s\nto\n%s.\n"
+ "Supported arguments: %s" % (
+ new_arg, text_input, text, str(args_v2)))
visitor = public_api.PublicAPIVisitor(conversion_visitor)
visitor.do_not_descend_map["tf"].append("contrib")
visitor.private_map["tf.compat"] = ["v1", "v2"]
traverse.traverse(tf.compat.v1, visitor)
+ def testReorderFileNeedsUpdate(self):
+ reordered_function_names = (
+ tf_upgrade_v2.TFAPIChangeSpec().reordered_function_names)
+ function_reorders = (
+ tf_upgrade_v2.TFAPIChangeSpec().function_reorders)
+
+ added_names_message = """Some function names in
+self.reordered_function_names are not in reorders_v2.py.
+Please run the following commands to update reorders_v2.py:
+bazel build tensorflow/tools/compatibility/update:generate_v2_reorders_map
+bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
+"""
+ removed_names_message = """%s in self.reorders_v2 does not match
+any name in self.reordered_function_names.
+Please run the following commands to update reorders_v2.py:
+bazel build tensorflow/tools/compatibility/update:generate_v2_reorders_map
+bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
+"""
+ self.assertTrue(
+ reordered_function_names.issubset(function_reorders),
+ added_names_message)
+ # function_reorders should contain reordered_function_names
+ # and their TensorFlow V1 aliases.
+ for name in function_reorders:
+ # get other names for this function
+ attr = get_symbol_for_name(tf.compat.v1, name)
+ _, attr = tf_decorator.unwrap(attr)
+ v1_names = tf_export.get_v1_names(attr)
+ self.assertTrue(v1_names)
+ v1_names = ["tf.%s" % n for n in v1_names]
+ # check if any other name is in
+ self.assertTrue(
+ any(n in reordered_function_names for n in v1_names),
+ removed_names_message % name)
+
def testRenameConstant(self):
text = "tf.MONOLITHIC_BUILD\n"
_, unused_report, unused_errors, new_text = self._upgrade(text)
@@ -336,6 +449,13 @@
)
self.assertEqual(new_text, expected_text)
+ def test_substr(self):
+ text = "tf.substr(input, pos, len, name, unit)\n"
+ _, unused_report, errors, new_text = self._upgrade(text)
+ self.assertEqual("tf.strings.substr(input=input, pos=pos, len=len, "
+ "name=name, unit=unit)\n", new_text)
+ self.assertEqual(errors, [])
+
def testColocateGradientsWithOps(self):
text = "tf.gradients(a, foo=False)\n"
_, unused_report, errors, new_text = self._upgrade(text)
@@ -451,9 +571,24 @@
text = "tf.nn.softmax_cross_entropy_with_logits_v2(labels, logits, dim=2)"
expected_text = (
"tf.nn.softmax_cross_entropy_with_logits(labels, logits, axis=2)")
- _, unused_report, unused_errors, new_text = self._upgrade(text)
+ _, unused_report, errors, new_text = self._upgrade(text)
self.assertEqual(new_text, expected_text)
+ self.assertFalse(errors)
+
+ def testSoftMaxCrossEntropyWithLogits(self):
+ text = "tf.nn.softmax_cross_entropy_with_logits(labels, logits, dim=2)"
+ expected_text = (
+ "tf.nn.softmax_cross_entropy_with_logits(labels, logits, dim=2)")
+ _, report, errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+ self.assertIn(
+ "tf.nn.softmax_cross_entropy_with_logits requires manual check.",
+ errors[0])
+ self.assertIn(
+ "tf.nn.softmax_cross_entropy_with_logits behavior has changed. ",
+ report)
+
def testSparseMatmul(self):
text = ("tf.sparse_matmul(a, b, c, d, e, f, g)\n")
expected_text = ("tf.linalg.matmul(a=a, b=b, transpose_a=c, transpose_b=d, "
@@ -465,7 +600,7 @@
text = "tf.nn.weighted_moments(x, axes, freq, name, kd)"
expected_text = (
"tf.nn.weighted_moments(x=x, axes=axes, frequency_weights=freq, "
- "name=name, keep_dims=kd)")
+ "name=name, keepdims=kd)")
_, unused_report, unused_errors, new_text = self._upgrade(text)
self.assertEqual(new_text, expected_text)
@@ -479,7 +614,7 @@
text = "tf.sparse.concat(ax, inp, name, exp, concat)"
expected_text = (
"tf.sparse.concat(axis=ax, sp_inputs=inp, name=name, "
- "expand_nonconcat_dim=exp, axis=concat)")
+ "expand_nonconcat_dims=exp, axis=concat)")
_, unused_report, unused_errors, new_text = self._upgrade(text)
self.assertEqual(new_text, expected_text)
@@ -492,6 +627,104 @@
_, unused_report, unused_errors, new_text = self._upgrade(text)
self.assertEqual(new_text, expected_text)
+ def testSpacetoBatch(self):
+ text = "tf.space_to_batch_nd(input, shape, paddings, name)"
+ expected_text = "tf.space_to_batch(input, shape, paddings, name)"
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ text = "tf.nn.space_to_batch(input, paddings, block_size, name)"
+ expected_text = (
+ "tf.space_to_batch(input=input, paddings=paddings, "
+ "block_shape=block_size, name=name)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testInTopK(self):
+ text = "tf.math.in_top_k(a, b, c, n)"
+ expected_text = (
+ "tf.math.in_top_k(predictions=a, targets=b, k=c, name=n)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testDepthToSpace(self):
+ text = "tf.nn.depth_to_space(input, block_size, name, data_format)"
+ expected_text = (
+ "tf.nn.depth_to_space(input=input, block_size=block_size, "
+ "name=name, data_format=data_format)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testEmbeddingLookup(self):
+ text = ("tf.nn.embedding_lookup(params, ids, partition_strategy, name, "
+ "validate_indices, max_norm)")
+ expected_text = ("tf.nn.embedding_lookup(params=params, ids=ids, "
+ "partition_strategy=partition_strategy, name=name, "
+ "validate_indices=validate_indices, max_norm=max_norm)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testEmbeddingLookupSparse(self):
+ text = ("tf.nn.embedding_lookup_sparse(params, sp_ids, sp_weights, "
+ "partition_strategy, name, combiner, max_norm)")
+ expected_text = ("tf.nn.embedding_lookup_sparse(params=params, "
+ "sp_ids=sp_ids, sp_weights=sp_weights, "
+ "partition_strategy=partition_strategy, name=name, "
+ "combiner=combiner, max_norm=max_norm)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testNnInTopK(self):
+ text = "tf.nn.in_top_k(predictions, targets, k, name)"
+ expected_text = ("tf.nn.in_top_k(predictions=predictions, "
+ "targets=targets, k=k, name=name)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testSpaceToDepth(self):
+ text = "tf.nn.space_to_depth(input, block_size, name, data_format)"
+ expected_text = ("tf.nn.space_to_depth(input=input, block_size=block_size, "
+ "name=name, data_format=data_format)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ def testPrint(self):
+ # tf.print() cannot be parsed unless we import print_function
+ text = """from __future__ import print_function
+tf.print()
+tf.print('abc')
+"""
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, text) # Text should stay the same
+
+ def testSparseSplit(self):
+ text = (
+ "tf.sparse_split(sp_input=sp_input, num_split=num_split, axis=axis, "
+ "name=name)")
+ expected_text = (
+ "tf.sparse.split(sp_input=sp_input, num_split=num_split, axis=axis, "
+ "name=name)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ text = (
+ "tf.sparse_split(sp_input=sp_input, num_split=num_split, "
+ "name=name, split_dim=axis)")
+ expected_text = (
+ "tf.sparse.split(sp_input=sp_input, num_split=num_split, "
+ "name=name, axis=axis)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
+ text = (
+ "tf.sparse.split(sp_input=sp_input, num_split=num_split, "
+ "name=name, split_dim=axis)")
+ expected_text = (
+ "tf.sparse.split(sp_input=sp_input, num_split=num_split, "
+ "name=name, axis=axis)")
+ _, unused_report, unused_errors, new_text = self._upgrade(text)
+ self.assertEqual(new_text, expected_text)
+
class TestUpgradeFiles(test_util.TensorFlowTestCase):
@@ -510,3 +743,4 @@
if __name__ == "__main__":
test_lib.main()
+
diff --git a/tensorflow/tools/compatibility/update/BUILD b/tensorflow/tools/compatibility/update/BUILD
index b9725a7..75bb0cf 100644
--- a/tensorflow/tools/compatibility/update/BUILD
+++ b/tensorflow/tools/compatibility/update/BUILD
@@ -15,3 +15,17 @@
"//tensorflow/tools/compatibility:tf_upgrade_v2_lib",
],
)
+
+py_binary(
+ name = "generate_v2_reorders_map",
+ srcs = ["generate_v2_reorders_map.py"],
+ srcs_version = "PY2AND3",
+ deps = [
+ "//tensorflow:tensorflow_py",
+ "//tensorflow/python:lib",
+ "//tensorflow/python:no_contrib",
+ "//tensorflow/tools/common:public_api",
+ "//tensorflow/tools/common:traverse",
+ "//tensorflow/tools/compatibility:tf_upgrade_v2_lib",
+ ],
+)
diff --git a/tensorflow/tools/compatibility/update/generate_v2_renames_map.py b/tensorflow/tools/compatibility/update/generate_v2_renames_map.py
index 19ad6c3..a2c5e7c 100644
--- a/tensorflow/tools/compatibility/update/generate_v2_renames_map.py
+++ b/tensorflow/tools/compatibility/update/generate_v2_renames_map.py
@@ -64,58 +64,6 @@
"""
-_TENSORFLOW_API_ATTR_V1 = (
- tf_export.API_ATTRS_V1[tf_export.TENSORFLOW_API_NAME].names)
-_TENSORFLOW_API_ATTR = tf_export.API_ATTRS[tf_export.TENSORFLOW_API_NAME].names
-_TENSORFLOW_CONSTANTS_ATTR_V1 = (
- tf_export.API_ATTRS_V1[tf_export.TENSORFLOW_API_NAME].constants)
-_TENSORFLOW_CONSTANTS_ATTR = (
- tf_export.API_ATTRS[tf_export.TENSORFLOW_API_NAME].constants)
-
-_ESTIMATOR_API_ATTR_V1 = (
- tf_export.API_ATTRS_V1[tf_export.ESTIMATOR_API_NAME].names)
-_ESTIMATOR_API_ATTR = tf_export.API_ATTRS[tf_export.ESTIMATOR_API_NAME].names
-_ESTIMATOR_CONSTANTS_ATTR_V1 = (
- tf_export.API_ATTRS_V1[tf_export.ESTIMATOR_API_NAME].constants)
-_ESTIMATOR_CONSTANTS_ATTR = (
- tf_export.API_ATTRS[tf_export.ESTIMATOR_API_NAME].constants)
-
-
-def get_v1_names(symbol):
- names_v1 = []
- if hasattr(symbol, _TENSORFLOW_API_ATTR_V1):
- names_v1.extend(getattr(symbol, _TENSORFLOW_API_ATTR_V1))
- if hasattr(symbol, _ESTIMATOR_API_ATTR_V1):
- names_v1.extend(getattr(symbol, _ESTIMATOR_API_ATTR_V1))
- return names_v1
-
-
-def get_v2_names(symbol):
- names_v2 = []
- if hasattr(symbol, _TENSORFLOW_API_ATTR):
- names_v2.extend(getattr(symbol, _TENSORFLOW_API_ATTR))
- if hasattr(symbol, _ESTIMATOR_API_ATTR):
- names_v2.extend(getattr(symbol, _ESTIMATOR_API_ATTR))
- return list(names_v2)
-
-
-def get_v1_constants(module):
- constants_v1 = []
- if hasattr(module, _TENSORFLOW_CONSTANTS_ATTR_V1):
- constants_v1.extend(getattr(module, _TENSORFLOW_CONSTANTS_ATTR_V1))
- if hasattr(module, _ESTIMATOR_CONSTANTS_ATTR_V1):
- constants_v1.extend(getattr(module, _ESTIMATOR_CONSTANTS_ATTR_V1))
- return constants_v1
-
-
-def get_v2_constants(module):
- constants_v2 = []
- if hasattr(module, _TENSORFLOW_CONSTANTS_ATTR):
- constants_v2.extend(getattr(module, _TENSORFLOW_CONSTANTS_ATTR))
- if hasattr(module, _ESTIMATOR_CONSTANTS_ATTR):
- constants_v2.extend(getattr(module, _ESTIMATOR_CONSTANTS_ATTR))
- return constants_v2
-
def get_canonical_name(v2_names, v1_name):
if v2_names:
@@ -131,7 +79,7 @@
"""Visitor that collects TF 2.0 names."""
for child in children:
_, attr = tf_decorator.unwrap(child[1])
- api_names_v2 = get_v2_names(attr)
+ api_names_v2 = tf_export.get_v2_names(attr)
for name in api_names_v2:
v2_names.add(name)
@@ -149,8 +97,8 @@
"""
renames = set()
for module in sys.modules.values():
- constants_v1_list = get_v1_constants(module)
- constants_v2_list = get_v2_constants(module)
+ constants_v1_list = tf_export.get_v1_constants(module)
+ constants_v2_list = tf_export.get_v2_constants(module)
# _tf_api_constants attribute contains a list of tuples:
# (api_names_list, constant_name)
@@ -186,8 +134,8 @@
"""Visitor that collects rename strings to add to rename_line_set."""
for child in children:
_, attr = tf_decorator.unwrap(child[1])
- api_names_v1 = get_v1_names(attr)
- api_names_v2 = get_v2_names(attr)
+ api_names_v1 = tf_export.get_v1_names(attr)
+ api_names_v2 = tf_export.get_v2_names(attr)
deprecated_api_names = set(api_names_v1) - set(api_names_v2)
for name in deprecated_api_names:
renames.add((name, get_canonical_name(api_names_v2, name)))
diff --git a/tensorflow/tools/compatibility/update/generate_v2_reorders_map.py b/tensorflow/tools/compatibility/update/generate_v2_reorders_map.py
new file mode 100644
index 0000000..0eb942d
--- /dev/null
+++ b/tensorflow/tools/compatibility/update/generate_v2_reorders_map.py
@@ -0,0 +1,132 @@
+# 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.
+# ==============================================================================
+# pylint: disable=line-too-long
+"""Script for updating tensorflow/tools/compatibility/reorders_v2.py.
+
+To update reorders_v2.py, run:
+ bazel build tensorflow/tools/compatibility/update:generate_v2_reorders_map
+ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
+"""
+# pylint: enable=line-too-long
+import tensorflow as tf
+
+# This import is needed so that TensorFlow python modules are in sys.modules.
+from tensorflow import python as tf_python # pylint: disable=unused-import
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import app
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_export
+from tensorflow.python.util import tf_inspect
+from tensorflow.tools.common import public_api
+from tensorflow.tools.common import traverse
+from tensorflow.tools.compatibility import tf_upgrade_v2
+
+
+_OUTPUT_FILE_PATH = 'third_party/tensorflow/tools/compatibility/reorders_v2.py'
+_FILE_HEADER = """# 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.
+# ==============================================================================
+# pylint: disable=line-too-long
+\"\"\"List of renames to apply when converting from TF 1.0 to TF 2.0.
+
+THIS FILE IS AUTOGENERATED: To update, please run:
+ bazel build tensorflow/tools/compatibility/update:generate_v2_reorders_map
+ bazel-bin/tensorflow/tools/compatibility/update/generate_v2_reorders_map
+This file should be updated whenever a function is added to
+self.reordered_function_names in tf_upgrade_v2.py.
+\"\"\"
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+"""
+
+
+def collect_function_arg_names(function_names):
+ """Determines argument names for reordered function signatures.
+
+ Args:
+ function_names: Functions to collect arguments for.
+
+ Returns:
+ Dictionary mapping function name to its arguments.
+ """
+ # Map from reordered function name to its arguments.
+ function_to_args = {}
+
+ def visit(unused_path, unused_parent, children):
+ """Visitor that collects arguments for reordered functions."""
+ for child in children:
+ _, attr = tf_decorator.unwrap(child[1])
+ api_names_v1 = tf_export.get_v1_names(attr)
+ api_names_v1 = ['tf.%s' % name for name in api_names_v1]
+ matches_function_names = any(
+ name in function_names for name in api_names_v1)
+ if matches_function_names:
+ arg_list = tf_inspect.getargspec(attr)[0]
+ for name in api_names_v1:
+ function_to_args[name] = arg_list
+
+ visitor = public_api.PublicAPIVisitor(visit)
+ visitor.do_not_descend_map['tf'].append('contrib')
+ visitor.do_not_descend_map['tf.compat'] = ['v1', 'v2']
+ traverse.traverse(tf, visitor)
+
+ return function_to_args
+
+
+def get_reorder_line(name, arg_list):
+ return ' \'%s\': %s' % (name, str(arg_list))
+
+
+def update_reorders_v2(output_file_path):
+ """Writes a Python dictionary mapping function name to argument order.
+
+ Args:
+ output_file_path: File path to write output to. Any existing contents
+ would be replaced.
+ """
+ reordered_function_names = (
+ tf_upgrade_v2.TFAPIChangeSpec().reordered_function_names)
+
+ all_reorders = collect_function_arg_names(reordered_function_names)
+
+ # List of reorder lines to write to output file in the form:
+ # 'tf.function_name': ['arg1', 'arg2', ...]
+ rename_lines = [
+ get_reorder_line(name, arg_names)
+ for name, arg_names in all_reorders.items()]
+ renames_file_text = '%sreorders = {\n%s\n}\n' % (
+ _FILE_HEADER, ',\n'.join(sorted(rename_lines)))
+ file_io.write_string_to_file(output_file_path, renames_file_text)
+
+
+def main(unused_argv):
+ update_reorders_v2(_OUTPUT_FILE_PATH)
+
+
+if __name__ == '__main__':
+ app.run(main=main)
diff --git a/tensorflow/tools/dockerfiles/assembler.py b/tensorflow/tools/dockerfiles/assembler.py
index bd8ed73..67a03202 100644
--- a/tensorflow/tools/dockerfiles/assembler.py
+++ b/tensorflow/tools/dockerfiles/assembler.py
@@ -170,6 +170,8 @@
type: string
dockerfile_exclusive_name:
type: string
+ dockerfile_subdirectory:
+ type: string
partials:
type: list
schema:
@@ -353,8 +355,9 @@
def find_first_slice_value(slices, key):
"""For a list of slices, get the first value for a certain key."""
for s in slices:
- if key in s:
+ if key in s and s[key] is not None:
return s[key]
+ return None
def assemble_tags(spec, cli_args, enabled_releases, all_partials):
@@ -389,6 +392,8 @@
used_partials = gather_slice_list_items(slices, 'partials')
used_tests = gather_slice_list_items(slices, 'tests')
test_runtime = find_first_slice_value(slices, 'test_runtime')
+ dockerfile_subdirectory = find_first_slice_value(
+ slices, 'dockerfile_subdirectory')
dockerfile_contents = merge_partials(spec['header'], used_partials,
all_partials)
@@ -398,6 +403,7 @@
'is_dockerfiles': release['is_dockerfiles'],
'upload_images': release['upload_images'],
'cli_args': tag_args,
+ 'dockerfile_subdirectory': dockerfile_subdirectory or '',
'partials': used_partials,
'tests': used_tests,
'test_runtime': test_runtime,
@@ -416,8 +422,7 @@
def upload_in_background(hub_repository, dock, image, tag):
"""Upload a docker image (to be used by multiprocessing)."""
image.tag(hub_repository, tag=tag)
- for line in list(dock.images.push(hub_repository, tag=tag, stream=True)):
- print(line)
+ print(dock.images.push(hub_repository, tag=tag))
def mkdir_p(path):
@@ -525,13 +530,15 @@
continue
# Write releases marked "is_dockerfiles" into the Dockerfile directory
- if FLAGS.construct_dockerfiles:
- path = os.path.join(FLAGS.dockerfile_dir, tag + '.Dockerfile')
- if tag_def['is_dockerfiles']:
- eprint('>> Writing {}...'.format(path))
- if not FLAGS.dry_run:
- with open(path, 'w') as f:
- f.write(tag_def['dockerfile_contents'])
+ if FLAGS.construct_dockerfiles and tag_def['is_dockerfiles']:
+ path = os.path.join(FLAGS.dockerfile_dir,
+ tag_def['dockerfile_subdirectory'],
+ tag + '.Dockerfile')
+ eprint('>> Writing {}...'.format(path))
+ if not FLAGS.dry_run:
+ mkdir_p(os.path.dirname(path))
+ with open(path, 'w') as f:
+ f.write(tag_def['dockerfile_contents'])
# Don't build any images for dockerfile-only releases
if not FLAGS.build_images:
diff --git a/tensorflow/tools/dockerfiles/dockerfiles/cpu-devel-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/devel-cpu-jupyter.Dockerfile
similarity index 94%
rename from tensorflow/tools/dockerfiles/dockerfiles/cpu-devel-jupyter.Dockerfile
rename to tensorflow/tools/dockerfiles/dockerfiles/devel-cpu-jupyter.Dockerfile
index 14ddf08..4326567 100644
--- a/tensorflow/tools/dockerfiles/dockerfiles/cpu-devel-jupyter.Dockerfile
+++ b/tensorflow/tools/dockerfiles/dockerfiles/devel-cpu-jupyter.Dockerfile
@@ -46,6 +46,9 @@
ENV CI_BUILD_PYTHON python
+# Check out TensorFlow source code if --build_arg CHECKOUT_TENSORFLOW=1
+ARG CHECKOUT_TF_SRC=0
+RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src
ARG USE_PYTHON_3_NOT_2
ARG _PY_SUFFIX=${USE_PYTHON_3_NOT_2:+3}
diff --git a/tensorflow/tools/dockerfiles/dockerfiles/cpu-devel.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/devel-cpu.Dockerfile
similarity index 92%
rename from tensorflow/tools/dockerfiles/dockerfiles/cpu-devel.Dockerfile
rename to tensorflow/tools/dockerfiles/dockerfiles/devel-cpu.Dockerfile
index 16973b4..5c5b2f9 100644
--- a/tensorflow/tools/dockerfiles/dockerfiles/cpu-devel.Dockerfile
+++ b/tensorflow/tools/dockerfiles/dockerfiles/devel-cpu.Dockerfile
@@ -46,6 +46,9 @@
ENV CI_BUILD_PYTHON python
+# Check out TensorFlow source code if --build_arg CHECKOUT_TENSORFLOW=1
+ARG CHECKOUT_TF_SRC=0
+RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src
ARG USE_PYTHON_3_NOT_2
ARG _PY_SUFFIX=${USE_PYTHON_3_NOT_2:+3}
diff --git a/tensorflow/tools/dockerfiles/dockerfiles/gpu-devel-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu-jupyter.Dockerfile
similarity index 95%
rename from tensorflow/tools/dockerfiles/dockerfiles/gpu-devel-jupyter.Dockerfile
rename to tensorflow/tools/dockerfiles/dockerfiles/devel-gpu-jupyter.Dockerfile
index 9ecaec3..8769e4e 100644
--- a/tensorflow/tools/dockerfiles/dockerfiles/gpu-devel-jupyter.Dockerfile
+++ b/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu-jupyter.Dockerfile
@@ -79,6 +79,10 @@
# NCCL 2.x
ENV TF_NCCL_VERSION=2
+# Check out TensorFlow source code if --build_arg CHECKOUT_TENSORFLOW=1
+ARG CHECKOUT_TF_SRC=0
+RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src
+
ARG USE_PYTHON_3_NOT_2
ARG _PY_SUFFIX=${USE_PYTHON_3_NOT_2:+3}
ARG PYTHON=python${_PY_SUFFIX}
diff --git a/tensorflow/tools/dockerfiles/dockerfiles/gpu-devel.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu.Dockerfile
similarity index 94%
rename from tensorflow/tools/dockerfiles/dockerfiles/gpu-devel.Dockerfile
rename to tensorflow/tools/dockerfiles/dockerfiles/devel-gpu.Dockerfile
index c79bc3c..809cda6 100644
--- a/tensorflow/tools/dockerfiles/dockerfiles/gpu-devel.Dockerfile
+++ b/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu.Dockerfile
@@ -79,6 +79,10 @@
# NCCL 2.x
ENV TF_NCCL_VERSION=2
+# Check out TensorFlow source code if --build_arg CHECKOUT_TENSORFLOW=1
+ARG CHECKOUT_TF_SRC=0
+RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src
+
ARG USE_PYTHON_3_NOT_2
ARG _PY_SUFFIX=${USE_PYTHON_3_NOT_2:+3}
ARG PYTHON=python${_PY_SUFFIX}
diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/cpu-devel.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/devel-cpu.partial.Dockerfile
similarity index 73%
rename from tensorflow/tools/dockerfiles/partials/ubuntu/cpu-devel.partial.Dockerfile
rename to tensorflow/tools/dockerfiles/partials/ubuntu/devel-cpu.partial.Dockerfile
index 901652c..a61dfbb 100644
--- a/tensorflow/tools/dockerfiles/partials/ubuntu/cpu-devel.partial.Dockerfile
+++ b/tensorflow/tools/dockerfiles/partials/ubuntu/devel-cpu.partial.Dockerfile
@@ -23,3 +23,6 @@
ENV CI_BUILD_PYTHON python
+# Check out TensorFlow source code if --build_arg CHECKOUT_TENSORFLOW=1
+ARG CHECKOUT_TF_SRC=0
+RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src
diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/nvidia-devel.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/devel-nvidia.partial.Dockerfile
similarity index 90%
rename from tensorflow/tools/dockerfiles/partials/ubuntu/nvidia-devel.partial.Dockerfile
rename to tensorflow/tools/dockerfiles/partials/ubuntu/devel-nvidia.partial.Dockerfile
index 48d457e..95f9875 100644
--- a/tensorflow/tools/dockerfiles/partials/ubuntu/nvidia-devel.partial.Dockerfile
+++ b/tensorflow/tools/dockerfiles/partials/ubuntu/devel-nvidia.partial.Dockerfile
@@ -55,3 +55,7 @@
# NCCL 2.x
ENV TF_NCCL_VERSION=2
+
+# Check out TensorFlow source code if --build_arg CHECKOUT_TENSORFLOW=1
+ARG CHECKOUT_TF_SRC=0
+RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src
diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml
index 4826ddd..19d96e7 100644
--- a/tensorflow/tools/dockerfiles/spec.yml
+++ b/tensorflow/tools/dockerfiles/spec.yml
@@ -33,21 +33,28 @@
# - nightly-py3
# - nightly-gpu (similar)
# - nightly-gpu-py3
-
+#
+# Releases are all treated differently by TensorFlow's CI systems.
releases:
+ # Built Nightly and pushed to tensorflow/tensorflow
nightly:
tag_specs:
- "{nightly}{py}{jupyter}"
+ - "{ubuntu-devel}{py}"
+ # Built per-release and pushed to tensorflow/tensorflow
+ # --arg _TAG_PREFIX=<val> should be set to "1.11" (for example) or "latest".
versioned:
tag_specs:
- "{_TAG_PREFIX}{ubuntu}{py}{jupyter}"
- ubuntu-dockerfiles:
+ # Dockerfiles stored in the TF repo; not pushed anywhere
+ dockerfiles:
is_dockerfiles: true
upload_images: false
tag_specs:
- "{ubuntu}{jupyter}"
+ - "{ubuntu-devel}{jupyter}"
slice_sets:
@@ -87,27 +94,33 @@
tests:
- import-gpu.sh
test_runtime: nvidia
- - add_to_name: "-devel"
- dockerfile_exclusive_name: "cpu-devel"
+
+ ubuntu-devel:
+ - add_to_name: "devel"
+ dockerfile_exclusive_name: "devel-cpu"
partials:
- ubuntu/version
- - ubuntu/cpu-devel
+ - ubuntu/devel-cpu
- ubuntu/python
- ubuntu/bazel
- shell
tests:
- build-cpu.sh
- - add_to_name: "-gpu-devel"
- dockerfile_exclusive_name: "gpu-devel"
+ args:
+ - CHECKOUT_TF_SRC=1
+ - add_to_name: "devel-gpu"
+ dockerfile_exclusive_name: "devel-gpu"
partials:
- ubuntu/version
- - ubuntu/nvidia-devel
+ - ubuntu/devel-nvidia
- ubuntu/python
- ubuntu/bazel
- shell
tests:
- build-gpu.sh
test_runtime: nvidia
+ args:
+ - CHECKOUT_TF_SRC=1
nightly:
- add_to_name: "nightly"
diff --git a/tensorflow/tools/graph_transforms/backports.cc b/tensorflow/tools/graph_transforms/backports.cc
index 5c153e8..041e7ee 100644
--- a/tensorflow/tools/graph_transforms/backports.cc
+++ b/tensorflow/tools/graph_transforms/backports.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/flatten_atrous.cc b/tensorflow/tools/graph_transforms/flatten_atrous.cc
index a6f7cb0..c80b28f 100644
--- a/tensorflow/tools/graph_transforms/flatten_atrous.cc
+++ b/tensorflow/tools/graph_transforms/flatten_atrous.cc
@@ -18,7 +18,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/fold_batch_norms.cc b/tensorflow/tools/graph_transforms/fold_batch_norms.cc
index 975b173..16a0f7d 100644
--- a/tensorflow/tools/graph_transforms/fold_batch_norms.cc
+++ b/tensorflow/tools/graph_transforms/fold_batch_norms.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/fold_constants_lib.cc b/tensorflow/tools/graph_transforms/fold_constants_lib.cc
index 6df2718..dcc36b1 100644
--- a/tensorflow/tools/graph_transforms/fold_constants_lib.cc
+++ b/tensorflow/tools/graph_transforms/fold_constants_lib.cc
@@ -33,7 +33,6 @@
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc b/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
index 156636a..fd546f8 100644
--- a/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
+++ b/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/fuse_convolutions.cc b/tensorflow/tools/graph_transforms/fuse_convolutions.cc
index df6e9e6..7754dde 100644
--- a/tensorflow/tools/graph_transforms/fuse_convolutions.cc
+++ b/tensorflow/tools/graph_transforms/fuse_convolutions.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/fuse_quantized_convolution.cc b/tensorflow/tools/graph_transforms/fuse_quantized_convolution.cc
index bd021d0..5aa2dd4 100644
--- a/tensorflow/tools/graph_transforms/fuse_quantized_convolution.cc
+++ b/tensorflow/tools/graph_transforms/fuse_quantized_convolution.cc
@@ -25,7 +25,6 @@
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
diff --git a/tensorflow/tools/graph_transforms/insert_logging.cc b/tensorflow/tools/graph_transforms/insert_logging.cc
index 3776654..ccc4854 100644
--- a/tensorflow/tools/graph_transforms/insert_logging.cc
+++ b/tensorflow/tools/graph_transforms/insert_logging.cc
@@ -22,7 +22,6 @@
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/obfuscate_names.cc b/tensorflow/tools/graph_transforms/obfuscate_names.cc
index c470b51..ee8ca3d 100644
--- a/tensorflow/tools/graph_transforms/obfuscate_names.cc
+++ b/tensorflow/tools/graph_transforms/obfuscate_names.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/quantize_nodes.cc b/tensorflow/tools/graph_transforms/quantize_nodes.cc
index a022f57..b139dad 100644
--- a/tensorflow/tools/graph_transforms/quantize_nodes.cc
+++ b/tensorflow/tools/graph_transforms/quantize_nodes.cc
@@ -23,7 +23,6 @@
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/quantize_weights.cc b/tensorflow/tools/graph_transforms/quantize_weights.cc
index cccae8a..a1a6e27 100644
--- a/tensorflow/tools/graph_transforms/quantize_weights.cc
+++ b/tensorflow/tools/graph_transforms/quantize_weights.cc
@@ -23,7 +23,6 @@
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/remove_attribute.cc b/tensorflow/tools/graph_transforms/remove_attribute.cc
index b1a04c0..0a76c2b 100644
--- a/tensorflow/tools/graph_transforms/remove_attribute.cc
+++ b/tensorflow/tools/graph_transforms/remove_attribute.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/remove_device.cc b/tensorflow/tools/graph_transforms/remove_device.cc
index 975fa37..fdd4316 100644
--- a/tensorflow/tools/graph_transforms/remove_device.cc
+++ b/tensorflow/tools/graph_transforms/remove_device.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/remove_nodes.cc b/tensorflow/tools/graph_transforms/remove_nodes.cc
index 05f036a..aa02886 100644
--- a/tensorflow/tools/graph_transforms/remove_nodes.cc
+++ b/tensorflow/tools/graph_transforms/remove_nodes.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/rename_attribute.cc b/tensorflow/tools/graph_transforms/rename_attribute.cc
index bd066aa..62897d4 100644
--- a/tensorflow/tools/graph_transforms/rename_attribute.cc
+++ b/tensorflow/tools/graph_transforms/rename_attribute.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/rename_op.cc b/tensorflow/tools/graph_transforms/rename_op.cc
index e1e13c1..9deee8b 100644
--- a/tensorflow/tools/graph_transforms/rename_op.cc
+++ b/tensorflow/tools/graph_transforms/rename_op.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/round_weights.cc b/tensorflow/tools/graph_transforms/round_weights.cc
index 72927e4..3a145ac 100644
--- a/tensorflow/tools/graph_transforms/round_weights.cc
+++ b/tensorflow/tools/graph_transforms/round_weights.cc
@@ -23,7 +23,6 @@
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/sort_by_execution_order.cc b/tensorflow/tools/graph_transforms/sort_by_execution_order.cc
index 43152d2..548f5ba 100644
--- a/tensorflow/tools/graph_transforms/sort_by_execution_order.cc
+++ b/tensorflow/tools/graph_transforms/sort_by_execution_order.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/graph_transforms/sparsify_gather.cc b/tensorflow/tools/graph_transforms/sparsify_gather.cc
index cc82100..bed51f8 100644
--- a/tensorflow/tools/graph_transforms/sparsify_gather.cc
+++ b/tensorflow/tools/graph_transforms/sparsify_gather.cc
@@ -25,7 +25,6 @@
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
diff --git a/tensorflow/tools/graph_transforms/strip_unused_nodes.cc b/tensorflow/tools/graph_transforms/strip_unused_nodes.cc
index ae9d0aa..d466f21 100644
--- a/tensorflow/tools/graph_transforms/strip_unused_nodes.cc
+++ b/tensorflow/tools/graph_transforms/strip_unused_nodes.cc
@@ -21,7 +21,6 @@
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
diff --git a/tensorflow/tools/pip_package/BUILD b/tensorflow/tools/pip_package/BUILD
index f34ac4c..4ed2f6c 100644
--- a/tensorflow/tools/pip_package/BUILD
+++ b/tensorflow/tools/pip_package/BUILD
@@ -18,6 +18,13 @@
"if_ngraph",
)
+# This flag specifies whether TensorFlow 2.0 API should be built instead
+# of 1.* API. Note that TensorFlow 2.0 API is currently under development.
+config_setting(
+ name = "api_version_2",
+ define_values = {"tf_api_version": "2"},
+)
+
# This returns a list of headers of all public header libraries (e.g.,
# framework, lib), and all of the transitive dependencies of those
# public headers. Not all of the headers returned by the filegroup
@@ -59,6 +66,42 @@
"setup.py",
":included_headers",
"//tensorflow:tensorflow_py",
+ "//tensorflow/examples/tutorials/mnist:package",
+ "//tensorflow/lite/python:interpreter_test_data",
+ "//tensorflow/lite/python:tflite_convert",
+ "//tensorflow/lite/toco/python:toco_from_protos",
+ # "//tensorflow/python/autograph/converters:converters",
+ # "//tensorflow/python/autograph/core:core",
+ "//tensorflow/python/autograph/core:test_lib",
+ # "//tensorflow/python/autograph/impl:impl",
+ # "//tensorflow/python/autograph/lang:lang",
+ # "//tensorflow/python/autograph/operators:operators",
+ # "//tensorflow/python/autograph/pyct:pyct",
+ # "//tensorflow/python/autograph/pyct/testing:testing",
+ # "//tensorflow/python/autograph/pyct/static_analysis:static_analysis",
+ "//tensorflow/python/autograph/pyct/common_transformers:common_transformers",
+ "//tensorflow/python:cond_v2",
+ "//tensorflow/python:distributed_framework_test_lib",
+ "//tensorflow/python:meta_graph_testdata",
+ "//tensorflow/python:spectral_ops_test_util",
+ "//tensorflow/python:util_example_parser_configuration",
+ "//tensorflow/python/data/experimental/kernel_tests/serialization:dataset_serialization_test_base",
+ "//tensorflow/python/data/experimental/kernel_tests:stats_dataset_test_base",
+ "//tensorflow/python/data/kernel_tests:test_base",
+ "//tensorflow/python/debug:debug_pip",
+ "//tensorflow/python/eager:eager_pip",
+ "//tensorflow/python/kernel_tests/signal:test_util",
+ "//tensorflow/python/kernel_tests/testdata:self_adjoint_eig_op_test_files",
+ "//tensorflow/python/ops/ragged:ragged_test_util",
+ "//tensorflow/python/saved_model:saved_model",
+ "//tensorflow/python/tools:tools_pip",
+ "//tensorflow/python/tools/api/generator:create_python_api",
+ "//tensorflow/python:test_ops",
+ "//tensorflow/python:while_v2",
+ "//tensorflow/tools/dist_test/server:grpc_tensorflow_server",
+]
+
+COMMON_PIP_DEPS_V1 = COMMON_PIP_DEPS + [
"//tensorflow/contrib/autograph:autograph",
"//tensorflow/contrib/boosted_trees:boosted_trees_pip",
"//tensorflow/contrib/cluster_resolver:cluster_resolver_pip",
@@ -87,38 +130,6 @@
"//tensorflow/contrib/tensor_forest/hybrid:hybrid_pip",
"//tensorflow/contrib/timeseries:timeseries_pip",
"//tensorflow/contrib/tpu",
- "//tensorflow/examples/tutorials/mnist:package",
- "//tensorflow/lite/python:interpreter_test_data",
- "//tensorflow/lite/python:tflite_convert",
- "//tensorflow/lite/toco/python:toco_from_protos",
- # "//tensorflow/python/autograph/converters:converters",
- # "//tensorflow/python/autograph/core:core",
- "//tensorflow/python/autograph/core:test_lib",
- # "//tensorflow/python/autograph/impl:impl",
- # "//tensorflow/python/autograph/lang:lang",
- # "//tensorflow/python/autograph/operators:operators",
- # "//tensorflow/python/autograph/pyct:pyct",
- # "//tensorflow/python/autograph/pyct/testing:testing",
- # "//tensorflow/python/autograph/pyct/static_analysis:static_analysis",
- "//tensorflow/python/autograph/pyct/common_transformers:common_transformers",
- "//tensorflow/python:cond_v2",
- "//tensorflow/python:distributed_framework_test_lib",
- "//tensorflow/python:meta_graph_testdata",
- "//tensorflow/python:spectral_ops_test_util",
- "//tensorflow/python:util_example_parser_configuration",
- "//tensorflow/python/data/experimental/kernel_tests/serialization:dataset_serialization_test_base",
- "//tensorflow/python/data/experimental/kernel_tests:stats_dataset_test_base",
- "//tensorflow/python/data/kernel_tests:test_base",
- "//tensorflow/python/debug:debug_pip",
- "//tensorflow/python/eager:eager_pip",
- "//tensorflow/python/kernel_tests/signal:test_util",
- "//tensorflow/python/kernel_tests/testdata:self_adjoint_eig_op_test_files",
- "//tensorflow/python/saved_model:saved_model",
- "//tensorflow/python/tools:tools_pip",
- "//tensorflow/python/tools/api/generator:create_python_api",
- "//tensorflow/python:test_ops",
- "//tensorflow/python:while_v2",
- "//tensorflow/tools/dist_test/server:grpc_tensorflow_server",
]
# On Windows, python binary is a zip file of runfiles tree.
@@ -127,7 +138,10 @@
py_binary(
name = "simple_console_for_windows",
srcs = ["simple_console_for_windows.py"],
- data = COMMON_PIP_DEPS + ["//tensorflow/python:pywrap_tensorflow_import_lib_file"],
+ data = select({
+ "api_version_2": COMMON_PIP_DEPS,
+ "//conditions:default": COMMON_PIP_DEPS_V1,
+ }) + ["//tensorflow/python:pywrap_tensorflow_import_lib_file"],
srcs_version = "PY2AND3",
deps = ["//tensorflow:tensorflow_py"],
)
@@ -231,7 +245,10 @@
"//tensorflow:windows": [
":simple_console_for_windows",
],
- "//conditions:default": COMMON_PIP_DEPS + [
+ "api_version_2": COMMON_PIP_DEPS + [
+ ":simple_console",
+ ],
+ "//conditions:default": COMMON_PIP_DEPS_V1 + [
":simple_console",
],
}) + if_mkl_ml(["//third_party/mkl:intel_binary_blob"]),
diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index 85c913f..a3da276 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -45,7 +45,7 @@
# This version string is semver compatible, but incompatible with pip.
# For pip, we will remove all '-' characters from this string, and use the
# result for pip.
-_VERSION = '1.12.0-rc0'
+_VERSION = '1.12.0'
REQUIRED_PACKAGES = [
'absl-py >= 0.1.6',
@@ -89,6 +89,8 @@
REQUIRED_PACKAGES[i] = 'tb-nightly >= 1.13.0a0, < 1.14.0a0'
if 'tensorflow_estimator' in pkg:
REQUIRED_PACKAGES[i] = 'tf-estimator-nightly'
+ if 'tensorflow_estimator' in pkg and '2.0' in project_name:
+ REQUIRED_PACKAGES[i] = 'tensorflow-estimator-2.0-preview'
# weakref.finalize and enum were introduced in Python 3.4
if sys.version_info < (3, 4):
diff --git a/tensorflow/version_check.bzl b/tensorflow/version_check.bzl
index 79e721d..74feaa1 100644
--- a/tensorflow/version_check.bzl
+++ b/tensorflow/version_check.bzl
@@ -1,48 +1,52 @@
""" Helpers to check minimum version of bazel."""
def _extract_version_number(bazel_version):
- """Extracts the semantic version number from a version string
+ """Extracts the semantic version number from a version string
- Args:
- bazel_version: the version string that begins with the semantic version
- e.g. "1.2.3rc1 abc1234" where "abc1234" is a commit hash.
+ Args:
+ bazel_version: the version string that begins with the semantic version
+ e.g. "1.2.3rc1 abc1234" where "abc1234" is a commit hash.
- Returns:
- The semantic version string, like "1.2.3".
- """
- for i in range(len(bazel_version)):
- c = bazel_version[i]
- if not (c.isdigit() or c == "."):
- return bazel_version[:i]
- return bazel_version
+ Returns:
+ The semantic version string, like "1.2.3".
+ """
+ for i in range(len(bazel_version)):
+ c = bazel_version[i]
+ if not (c.isdigit() or c == "."):
+ return bazel_version[:i]
+ return bazel_version
# Parse the bazel version string from `native.bazel_version`.
# e.g.
# "0.10.0rc1 abc123d" => (0, 10, 0)
# "0.3.0" => (0, 3, 0)
def _parse_bazel_version(bazel_version):
- """Parses a version string into a 3-tuple of ints
+ """Parses a version string into a 3-tuple of ints
- int tuples can be compared directly using binary operators (<, >).
+ int tuples can be compared directly using binary operators (<, >).
- Args:
- bazel_version: the Bazel version string
+ Args:
+ bazel_version: the Bazel version string
- Returns:
- An int 3-tuple of a (major, minor, patch) version.
- """
+ Returns:
+ An int 3-tuple of a (major, minor, patch) version.
+ """
- version = _extract_version_number(bazel_version)
- return tuple([int(n) for n in version.split(".")])
+ version = _extract_version_number(bazel_version)
+ return tuple([int(n) for n in version.split(".")])
def check_bazel_version_at_least(minimum_bazel_version):
- if "bazel_version" not in dir(native):
- fail("\nCurrent Bazel version is lower than 0.2.1, expected at least %s\n" % minimum_bazel_version)
- elif not native.bazel_version:
- print("\nCurrent Bazel is not a release version, cannot check for compatibility.")
- print("Make sure that you are running at least Bazel %s.\n" % minimum_bazel_version)
- return
+ if "bazel_version" not in dir(native):
+ fail("\nCurrent Bazel version is lower than 0.2.1, expected at least %s\n" % minimum_bazel_version)
+ elif not native.bazel_version:
+ print("\nCurrent Bazel is not a release version, cannot check for compatibility.")
+ print("Make sure that you are running at least Bazel %s.\n" % minimum_bazel_version)
+ return
- if _parse_bazel_version(native.bazel_version) < _parse_bazel_version(minimum_bazel_version):
- fail("\nCurrent Bazel version is {}, expected at least {}\n".format(
- native.bazel_version, minimum_bazel_version))
+ if _parse_bazel_version(native.bazel_version) < _parse_bazel_version(minimum_bazel_version):
+ fail("\nCurrent Bazel version is {}, expected at least {}\n".format(
+ native.bazel_version,
+ minimum_bazel_version,
+ ))
+
+parse_bazel_version = _parse_bazel_version
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 2074f632..b9ae329 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -123,22 +123,22 @@
tf_http_archive(
name = "com_google_absl",
build_file = clean_dep("//third_party:com_google_absl.BUILD"),
- sha256 = "3cf6132129ba87f0781c383bfaf381b7174b5818e81fffcc5d04bb451154f0f2",
- strip_prefix = "abseil-cpp-f95179062eb65ce40895cc76f1398cce25394369",
+ sha256 = "c2f8a1a399994df49db348a4725933b12fc807909cee21d48e46f53a28e79d4b",
+ strip_prefix = "abseil-cpp-8fbcdb90952c57828c4a9c2f6d79fcd7cae9088f",
urls = [
- "https://mirror.bazel.build/github.com/abseil/abseil-cpp/archive/f95179062eb65ce40895cc76f1398cce25394369.tar.gz",
- "https://github.com/abseil/abseil-cpp/archive/f95179062eb65ce40895cc76f1398cce25394369.tar.gz",
+ "https://mirror.bazel.build/github.com/abseil/abseil-cpp/archive/8fbcdb90952c57828c4a9c2f6d79fcd7cae9088f.tar.gz",
+ "https://github.com/abseil/abseil-cpp/archive/8fbcdb90952c57828c4a9c2f6d79fcd7cae9088f.tar.gz",
],
)
tf_http_archive(
name = "eigen_archive",
build_file = clean_dep("//third_party:eigen.BUILD"),
- sha256 = "37a483ec219c43219b6e0fc07e799277a4a36abb2b9f4162cfcd256aa211eae8",
- strip_prefix = "eigen-eigen-2e50f4a5542a",
+ sha256 = "753fbb58d0a49b6bcbcfb126ebfa2e21fc97f7471529ba835a096008ce588d8a",
+ strip_prefix = "eigen-eigen-9f48e814419e",
urls = [
- "https://mirror.bazel.build/bitbucket.org/eigen/eigen/get/2e50f4a5542a.tar.gz",
- "https://bitbucket.org/eigen/eigen/get/2e50f4a5542a.tar.gz",
+ "https://mirror.bazel.build/bitbucket.org/eigen/eigen/get/9f48e814419e.tar.gz",
+ "https://bitbucket.org/eigen/eigen/get/9f48e814419e.tar.gz",
],
)
@@ -168,26 +168,26 @@
tf_http_archive(
name = "com_googlesource_code_re2",
- sha256 = "803c7811146edeef8f91064de37c6f19136ff01a2a8cdb3230e940b2fd9f07fe",
- strip_prefix = "re2-2018-07-01",
+ sha256 = "a31397714a353587413d307337d0b58f8a2e20e2b9d02f2e24e3463fa4eeda81",
+ strip_prefix = "re2-2018-10-01",
system_build_file = clean_dep("//third_party/systemlibs:re2.BUILD"),
urls = [
- "https://mirror.bazel.build/github.com/google/re2/archive/2018-07-01.tar.gz",
- "https://github.com/google/re2/archive/2018-07-01.tar.gz",
+ "https://mirror.bazel.build/github.com/google/re2/archive/2018-10-01.tar.gz",
+ "https://github.com/google/re2/archive/2018-10-01.tar.gz",
],
)
tf_http_archive(
name = "com_github_googlecloudplatform_google_cloud_cpp",
- sha256 = "3ade2072e6588ff56c0434abe6c63aa5f3f2d56be15a299bafc7e9cdf0a12c17",
- strip_prefix = "google-cloud-cpp-0.3.0",
+ sha256 = "44eee8bd47cbd5ff192e895b45f9f913e2e117f10fdb9af0fd3b1a87a7b53bc3",
+ strip_prefix = "google-cloud-cpp-0.4.0",
system_build_file = clean_dep("//third_party/systemlibs:google_cloud_cpp.BUILD"),
system_link_files = {
"//third_party/systemlibs:google_cloud_cpp.google.cloud.bigtable.BUILD": "google/cloud/bigtable/BUILD",
},
urls = [
- "https://mirror.bazel.build/github.com/GoogleCloudPlatform/google-cloud-cpp/archive/v0.3.0.tar.gz",
- "https://github.com/GoogleCloudPlatform/google-cloud-cpp/archive/v0.3.0.tar.gz",
+ "https://mirror.bazel.build/github.com/GoogleCloudPlatform/google-cloud-cpp/archive/v0.4.0.tar.gz",
+ "https://github.com/GoogleCloudPlatform/google-cloud-cpp/archive/v0.4.0.tar.gz",
],
)
@@ -347,16 +347,20 @@
)
PROTOBUF_URLS = [
- "https://mirror.bazel.build/github.com/google/protobuf/archive/v3.6.1.1.tar.gz",
- "https://github.com/google/protobuf/archive/v3.6.1.1.tar.gz",
+ "https://mirror.bazel.build/github.com/protocolbuffers/protobuf/archive/v3.6.1.2.tar.gz",
+ "https://github.com/protocolbuffers/protobuf/archive/v3.6.1.2.tar.gz",
]
- PROTOBUF_SHA256 = "1ade182f91f0fa6c6116195def5d22270e01b9d03fe91319e4c6215022d0d24b"
- PROTOBUF_STRIP_PREFIX = "protobuf-3.6.1.1"
+ PROTOBUF_SHA256 = "2244b0308846bb22b4ff0bcc675e99290ff9f1115553ae9671eba1030af31bc0"
+ PROTOBUF_STRIP_PREFIX = "protobuf-3.6.1.2"
tf_http_archive(
name = "protobuf_archive",
sha256 = PROTOBUF_SHA256,
strip_prefix = PROTOBUF_STRIP_PREFIX,
+ system_build_file = clean_dep("//third_party/systemlibs:protobuf.BUILD"),
+ system_link_files = {
+ "//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
+ },
urls = PROTOBUF_URLS,
)
@@ -367,6 +371,10 @@
name = "com_google_protobuf",
sha256 = PROTOBUF_SHA256,
strip_prefix = PROTOBUF_STRIP_PREFIX,
+ system_build_file = clean_dep("//third_party/systemlibs:protobuf.BUILD"),
+ system_link_files = {
+ "//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
+ },
urls = PROTOBUF_URLS,
)
@@ -374,6 +382,10 @@
name = "com_google_protobuf_cc",
sha256 = PROTOBUF_SHA256,
strip_prefix = PROTOBUF_STRIP_PREFIX,
+ system_build_file = clean_dep("//third_party/systemlibs:protobuf.BUILD"),
+ system_link_files = {
+ "//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
+ },
urls = PROTOBUF_URLS,
)
@@ -484,11 +496,11 @@
tf_http_archive(
name = "llvm",
build_file = clean_dep("//third_party/llvm:llvm.autogenerated.BUILD"),
- sha256 = "34170a4aa07e434dd537d98a705dcf1b3901f73820fe1d6b9370e8c1c94e9157",
- strip_prefix = "llvm-0487bd8f42c8b38166ff825d56014d0ff49db604",
+ sha256 = "55769c91b9f5b5255d58a1ecd88e690a4e192dc8cbdf8f984596649abe3b5433",
+ strip_prefix = "llvm-2ba3294845dedcbb27dc49287bfbcdb49aa1e6b7",
urls = [
- "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/0487bd8f42c8b38166ff825d56014d0ff49db604.tar.gz",
- "https://github.com/llvm-mirror/llvm/archive/0487bd8f42c8b38166ff825d56014d0ff49db604.tar.gz",
+ "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/2ba3294845dedcbb27dc49287bfbcdb49aa1e6b7.tar.gz",
+ "https://github.com/llvm-mirror/llvm/archive/2ba3294845dedcbb27dc49287bfbcdb49aa1e6b7.tar.gz",
],
)
@@ -722,12 +734,22 @@
)
tf_http_archive(
- name = "tflite_mobilenet",
- build_file = clean_dep("//third_party:tflite_mobilenet.BUILD"),
- sha256 = "23f814d1c076bdf03715dfb6cab3713aa4fbdf040fd5448c43196bd2e97a4c1b",
+ name = "tflite_mobilenet_float",
+ build_file = clean_dep("//third_party:tflite_mobilenet_float.BUILD"),
+ sha256 = "2fadeabb9968ec6833bee903900dda6e61b3947200535874ce2fe42a8493abc0",
urls = [
- "https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip",
- "https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip",
+ "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224.tgz",
+ "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224.tgz",
+ ],
+ )
+
+ tf_http_archive(
+ name = "tflite_mobilenet_quant",
+ build_file = clean_dep("//third_party:tflite_mobilenet_quant.BUILD"),
+ sha256 = "d32432d28673a936b2d6281ab0600c71cf7226dfe4cdcef3012555f691744166",
+ urls = [
+ "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
+ "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
],
)
diff --git a/third_party/gpus/crosstool/BUILD.tpl b/third_party/gpus/crosstool/BUILD.tpl
index c8812fa..1260b26 100644
--- a/third_party/gpus/crosstool/BUILD.tpl
+++ b/third_party/gpus/crosstool/BUILD.tpl
@@ -22,6 +22,7 @@
"local|compiler": ":cc-compiler-local",
"darwin|compiler": ":cc-compiler-darwin",
"x64_windows|msvc-cl": ":cc-compiler-windows",
+ "x64_windows": ":cc-compiler-windows",
},
)
@@ -41,6 +42,7 @@
# last on the command line and contain all shared libraries to link, so all
# regular options will be left of them.
supports_param_files = 1,
+ toolchain_identifier = "local_linux",
)
cc_toolchain(
@@ -55,6 +57,7 @@
static_runtime_libs = [":empty"],
strip_files = ":empty",
supports_param_files = 0,
+ toolchain_identifier = "local_darwin",
)
cc_toolchain(
@@ -69,6 +72,7 @@
static_runtime_libs = [":empty"],
strip_files = ":empty",
supports_param_files = 1,
+ toolchain_identifier = "local_windows",
)
filegroup(
diff --git a/third_party/gpus/crosstool/CROSSTOOL.tpl b/third_party/gpus/crosstool/CROSSTOOL.tpl
index 921188c..5ca9b2d 100644
--- a/third_party/gpus/crosstool/CROSSTOOL.tpl
+++ b/third_party/gpus/crosstool/CROSSTOOL.tpl
@@ -642,6 +642,31 @@
name: "no_legacy_features"
}
+ # TODO(klimek): Previously we were using a .bat file to start python to run
+ # the python script that can redirect to nvcc - unfortunately .bat files
+ # have a rather short maximum length for command lines (8k). Instead, we
+ # now use the python binary as the compiler and pass the python script to
+ # it at the start of the command line. Investigate different possibilities
+ # to run the nvcc wrapper, either using pyinstaller --onefile, or writing
+ # a small C++ wrapper to redirect.
+ feature {
+ name: "redirector"
+ enabled: true
+ flag_set {
+ action: "c-compile"
+ action: "c++-compile"
+ action: "c++-module-compile"
+ action: "c++-module-codegen"
+ action: "c++-header-parsing"
+ action: "assemble"
+ action: "preprocess-assemble"
+ flag_group {
+ flag: "-B"
+ flag: "external/local_config_cuda/crosstool/windows/msvc_wrapper_for_nvcc.py"
+ }
+ }
+ }
+
# Suppress startup banner.
feature {
name: "nologo"
diff --git a/third_party/gpus/crosstool/windows/msvc_wrapper_for_nvcc.bat.tpl b/third_party/gpus/crosstool/windows/msvc_wrapper_for_nvcc.bat.tpl
deleted file mode 100644
index 8f8fb3e..0000000
--- a/third_party/gpus/crosstool/windows/msvc_wrapper_for_nvcc.bat.tpl
+++ /dev/null
@@ -1,20 +0,0 @@
-:: 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.
-:: =============================================================================
-
-:: Invoke msvc_wrapper_for_nvcc.py, which is located in the same directory.
-@echo OFF
-set arg0=%~0
-for %%F in ("%arg0%") do set DRIVER_BIN=%%~dpF
-"%{python_binary}" -B "%DRIVER_BIN%\msvc_wrapper_for_nvcc.py" %*
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index 03c67bc..8aa5b89 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -190,7 +190,7 @@
get_env_var(repository_ctx, "TMP", "C:\\Windows\\Temp").replace(
"\\", "\\\\"),)
- msvc_cl_path = "windows/msvc_wrapper_for_nvcc.bat"
+ msvc_cl_path = _get_python_bin(repository_ctx)
msvc_ml_path = find_msvc_tool(repository_ctx, vc_path, "ml64.exe").replace(
"\\", "/")
msvc_link_path = find_msvc_tool(repository_ctx, vc_path, "link.exe").replace(
@@ -1426,7 +1426,6 @@
repository_ctx.file(
"crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc", "")
repository_ctx.file("crosstool/windows/msvc_wrapper_for_nvcc.py", "")
- repository_ctx.file("crosstool/windows/msvc_wrapper_for_nvcc.bat", "")
else:
cuda_defines[
"%{host_compiler_path}"] = "clang/bin/crosstool_wrapper_driver_is_not_gcc"
@@ -1486,13 +1485,6 @@
"crosstool:windows/msvc_wrapper_for_nvcc.py",
wrapper_defines,
)
- _tpl(
- repository_ctx,
- "crosstool:windows/msvc_wrapper_for_nvcc.bat",
- {
- "%{python_binary}": _get_python_bin(repository_ctx),
- },
- )
_tpl(
repository_ctx,
diff --git a/third_party/icu/BUILD.system b/third_party/icu/BUILD.system
index 328e412..8a88a6e 100644
--- a/third_party/icu/BUILD.system
+++ b/third_party/icu/BUILD.system
@@ -1,13 +1,19 @@
+package(
+ default_visibility = ["//visibility:public"],
+)
+
licenses(["notice"]) # Apache 2.0
filegroup(
name = "icu4c/LICENSE",
- visibility = ["//visibility:public"],
)
filegroup(
name = "icu4j/main/shared/licenses/LICENSE",
- visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "headers",
)
cc_library(
@@ -15,7 +21,6 @@
deps = [
":icuuc",
],
- visibility = ["//visibility:public"],
)
cc_library(
diff --git a/third_party/keras_applications_archive/BUILD.system b/third_party/keras_applications_archive/BUILD.system
new file mode 100644
index 0000000..a3b58f1
--- /dev/null
+++ b/third_party/keras_applications_archive/BUILD.system
@@ -0,0 +1,13 @@
+# Description: Keras Applications: set of pre-trained deep learning models.
+
+licenses(["notice"]) # MIT
+
+filegroup(
+ name = "LICENSE",
+ visibility = ["//visibility:public"],
+)
+
+py_library(
+ name = "keras_applications",
+ visibility = ["//visibility:public"],
+)
diff --git a/third_party/keras_applications_archive/workspace.bzl b/third_party/keras_applications_archive/workspace.bzl
index e90630f..cf9d15c 100644
--- a/third_party/keras_applications_archive/workspace.bzl
+++ b/third_party/keras_applications_archive/workspace.bzl
@@ -12,4 +12,5 @@
"https://github.com/keras-team/keras-applications/archive/1.0.6.tar.gz",
],
build_file = "//third_party/keras_applications_archive:BUILD.bazel",
+ system_build_file = "//third_party/keras_applications_archive:BUILD.system",
)
diff --git a/third_party/llvm/llvm.autogenerated.BUILD b/third_party/llvm/llvm.autogenerated.BUILD
index eb468aa..6599b9e 100644
--- a/third_party/llvm/llvm.autogenerated.BUILD
+++ b/third_party/llvm/llvm.autogenerated.BUILD
@@ -2241,7 +2241,6 @@
deps = [
":code_gen",
":config",
- ":core",
":support",
],
)
diff --git a/third_party/mkl_dnn/mkldnn.BUILD b/third_party/mkl_dnn/mkldnn.BUILD
index 7a8ed3b..d80c713 100644
--- a/third_party/mkl_dnn/mkldnn.BUILD
+++ b/third_party/mkl_dnn/mkldnn.BUILD
@@ -63,3 +63,27 @@
"//conditions:default": [],
}),
)
+
+cc_library(
+ name = "mkldnn_single_threaded",
+ srcs = glob([
+ "src/common/*.cpp",
+ "src/cpu/*.cpp",
+ "src/cpu/gemm/*.cpp",
+ ]),
+ hdrs = glob(["include/*"]),
+ copts = [
+ "-fexceptions",
+ "-DMKLDNN_THR=MKLDNN_THR_SEQ", # Disables threading.
+ ],
+ includes = [
+ "include",
+ "src",
+ "src/common",
+ "src/cpu",
+ "src/cpu/gemm",
+ "src/cpu/xbyak",
+ ],
+ nocopts = "-fno-exceptions",
+ visibility = ["//visibility:public"],
+)
diff --git a/third_party/mpi/mpi.bzl b/third_party/mpi/mpi.bzl
index 38ce91c..3a48335 100644
--- a/third_party/mpi/mpi.bzl
+++ b/third_party/mpi/mpi.bzl
@@ -2,16 +2,16 @@
#based on the configuration options return one or the other
def mpi_hdr():
- MPI_LIB_IS_OPENMPI=True
- hdrs = []
+ MPI_LIB_IS_OPENMPI = True
+ hdrs = []
if MPI_LIB_IS_OPENMPI:
- hdrs = ["mpi.h", "mpi_portable_platform.h"] #When using OpenMPI
+ hdrs = ["mpi.h", "mpi_portable_platform.h"] #When using OpenMPI
else:
- hdrs = ["mpi.h", "mpio.h", "mpicxx.h"] #When using MVAPICH
+ hdrs = ["mpi.h", "mpio.h", "mpicxx.h"] #When using MVAPICH
return hdrs
def if_mpi(if_true, if_false = []):
return select({
"//tensorflow:with_mpi_support": if_true,
- "//conditions:default": if_false
+ "//conditions:default": if_false,
})
diff --git a/third_party/systemlibs/protobuf.BUILD b/third_party/systemlibs/protobuf.BUILD
new file mode 100644
index 0000000..4b1cf39
--- /dev/null
+++ b/third_party/systemlibs/protobuf.BUILD
@@ -0,0 +1,104 @@
+load(
+ "@protobuf_archive//:protobuf.bzl",
+ "proto_gen",
+ "py_proto_library",
+ "cc_proto_library",
+)
+
+licenses(["notice"])
+
+filegroup(
+ name = "LICENSE",
+ visibility = ["//visibility:public"],
+)
+
+HEADERS = [
+ "google/protobuf/any.pb.h",
+ "google/protobuf/any.proto",
+ "google/protobuf/arena.h",
+ "google/protobuf/compiler/importer.h",
+ "google/protobuf/descriptor.h",
+ "google/protobuf/descriptor.pb.h",
+ "google/protobuf/descriptor.proto",
+ "google/protobuf/duration.pb.h",
+ "google/protobuf/duration.proto",
+ "google/protobuf/dynamic_message.h",
+ "google/protobuf/empty.pb.h",
+ "google/protobuf/empty.proto",
+ "google/protobuf/field_mask.pb.h",
+ "google/protobuf/field_mask.proto",
+ "google/protobuf/io/coded_stream.h",
+ "google/protobuf/io/zero_copy_stream.h",
+ "google/protobuf/io/zero_copy_stream_impl_lite.h",
+ "google/protobuf/map.h",
+ "google/protobuf/repeated_field.h",
+ "google/protobuf/text_format.h",
+ "google/protobuf/timestamp.pb.h",
+ "google/protobuf/timestamp.proto",
+ "google/protobuf/util/json_util.h",
+ "google/protobuf/util/type_resolver_util.h",
+ "google/protobuf/wrappers.pb.h",
+ "google/protobuf/wrappers.proto",
+]
+
+genrule(
+ name = "link_headers",
+ outs = HEADERS,
+ cmd = """
+ for i in $(OUTS); do
+ f=$${i#$(@D)/}
+ mkdir -p $(@D)/$${f%/*}
+ ln -sf $(INCLUDEDIR)/$$f $(@D)/$$f
+ done
+ """,
+)
+
+cc_library(
+ name = "protobuf",
+ hdrs = HEADERS,
+ linkopts = ["-lprotobuf"],
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "protobuf_headers",
+ hdrs = HEADERS,
+ linkopts = ["-lprotobuf"],
+ visibility = ["//visibility:public"],
+)
+
+cc_library(
+ name = "protoc_lib",
+ linkopts = ["-lprotoc"],
+ visibility = ["//visibility:public"],
+)
+
+genrule(
+ name = "protoc",
+ outs = ["protoc.bin"],
+ cmd = "ln -s $$(which protoc) $@",
+ executable = 1,
+ visibility = ["//visibility:public"],
+)
+
+cc_proto_library(
+ name = "cc_wkt_protos",
+ hdrs = HEADERS,
+ internal_bootstrap_hack = 1,
+ protoc = ":protoc",
+ visibility = ["//visibility:public"],
+)
+
+proto_gen(
+ name = "protobuf_python_genproto",
+ includes = ["."],
+ protoc = "@protobuf_archive//:protoc",
+ visibility = ["//visibility:public"],
+)
+
+py_library(
+ name = "protobuf_python",
+ data = [":link_headers"],
+ srcs_version = "PY2AND3",
+ visibility = ["//visibility:public"],
+)
diff --git a/third_party/systemlibs/protobuf.bzl b/third_party/systemlibs/protobuf.bzl
new file mode 100644
index 0000000..2aa7561
--- /dev/null
+++ b/third_party/systemlibs/protobuf.bzl
@@ -0,0 +1,425 @@
+def _GetPath(ctx, path):
+ if ctx.label.workspace_root:
+ return ctx.label.workspace_root + "/" + path
+ else:
+ return path
+
+def _IsNewExternal(ctx):
+ # Bazel 0.4.4 and older have genfiles paths that look like:
+ # bazel-out/local-fastbuild/genfiles/external/repo/foo
+ # After the exec root rearrangement, they look like:
+ # ../repo/bazel-out/local-fastbuild/genfiles/foo
+ return ctx.label.workspace_root.startswith("../")
+
+def _GenDir(ctx):
+ if _IsNewExternal(ctx):
+ # We are using the fact that Bazel 0.4.4+ provides repository-relative paths
+ # for ctx.genfiles_dir.
+ return ctx.genfiles_dir.path + (
+ "/" + ctx.attr.includes[0] if ctx.attr.includes and ctx.attr.includes[0] else ""
+ )
+
+ # This means that we're either in the old version OR the new version in the local repo.
+ # Either way, appending the source path to the genfiles dir works.
+ return ctx.var["GENDIR"] + "/" + _SourceDir(ctx)
+
+def _SourceDir(ctx):
+ if not ctx.attr.includes:
+ return ctx.label.workspace_root
+ if not ctx.attr.includes[0]:
+ return _GetPath(ctx, ctx.label.package)
+ if not ctx.label.package:
+ return _GetPath(ctx, ctx.attr.includes[0])
+ return _GetPath(ctx, ctx.label.package + "/" + ctx.attr.includes[0])
+
+def _CcHdrs(srcs, use_grpc_plugin = False):
+ ret = [s[:-len(".proto")] + ".pb.h" for s in srcs]
+ if use_grpc_plugin:
+ ret += [s[:-len(".proto")] + ".grpc.pb.h" for s in srcs]
+ return ret
+
+def _CcSrcs(srcs, use_grpc_plugin = False):
+ ret = [s[:-len(".proto")] + ".pb.cc" for s in srcs]
+ if use_grpc_plugin:
+ ret += [s[:-len(".proto")] + ".grpc.pb.cc" for s in srcs]
+ return ret
+
+def _CcOuts(srcs, use_grpc_plugin = False):
+ return _CcHdrs(srcs, use_grpc_plugin) + _CcSrcs(srcs, use_grpc_plugin)
+
+def _PyOuts(srcs, use_grpc_plugin = False):
+ ret = [s[:-len(".proto")] + "_pb2.py" for s in srcs]
+ if use_grpc_plugin:
+ ret += [s[:-len(".proto")] + "_pb2_grpc.py" for s in srcs]
+ return ret
+
+def _RelativeOutputPath(path, include, dest = ""):
+ if include == None:
+ return path
+
+ if not path.startswith(include):
+ fail("Include path %s isn't part of the path %s." % (include, path))
+
+ if include and include[-1] != "/":
+ include = include + "/"
+ if dest and dest[-1] != "/":
+ dest = dest + "/"
+
+ path = path[len(include):]
+ return dest + path
+
+def _proto_gen_impl(ctx):
+ """General implementation for generating protos"""
+ srcs = ctx.files.srcs
+ deps = []
+ deps += ctx.files.srcs
+ source_dir = _SourceDir(ctx)
+ gen_dir = _GenDir(ctx)
+ if source_dir:
+ import_flags = ["-I" + source_dir, "-I" + gen_dir]
+ else:
+ import_flags = ["-I."]
+
+ for dep in ctx.attr.deps:
+ import_flags += dep.proto.import_flags
+ deps += dep.proto.deps
+
+ args = []
+ if ctx.attr.gen_cc:
+ args += ["--cpp_out=" + gen_dir]
+ if ctx.attr.gen_py:
+ args += ["--python_out=" + gen_dir]
+
+ inputs = srcs + deps
+ if ctx.executable.plugin:
+ plugin = ctx.executable.plugin
+ lang = ctx.attr.plugin_language
+ if not lang and plugin.basename.startswith("protoc-gen-"):
+ lang = plugin.basename[len("protoc-gen-"):]
+ if not lang:
+ fail("cannot infer the target language of plugin", "plugin_language")
+
+ outdir = gen_dir
+ if ctx.attr.plugin_options:
+ outdir = ",".join(ctx.attr.plugin_options) + ":" + outdir
+ args += ["--plugin=protoc-gen-%s=%s" % (lang, plugin.path)]
+ args += ["--%s_out=%s" % (lang, outdir)]
+ inputs += [plugin]
+
+ if args:
+ ctx.action(
+ inputs = inputs,
+ outputs = ctx.outputs.outs,
+ arguments = args + import_flags + [s.path for s in srcs],
+ executable = ctx.executable.protoc,
+ mnemonic = "ProtoCompile",
+ use_default_shell_env = True,
+ )
+
+ return struct(
+ proto = struct(
+ srcs = srcs,
+ import_flags = import_flags,
+ deps = deps,
+ ),
+ )
+
+proto_gen = rule(
+ attrs = {
+ "srcs": attr.label_list(allow_files = True),
+ "deps": attr.label_list(providers = ["proto"]),
+ "includes": attr.string_list(),
+ "protoc": attr.label(
+ cfg = "host",
+ executable = True,
+ single_file = True,
+ mandatory = True,
+ ),
+ "plugin": attr.label(
+ cfg = "host",
+ allow_files = True,
+ executable = True,
+ ),
+ "plugin_language": attr.string(),
+ "plugin_options": attr.string_list(),
+ "gen_cc": attr.bool(),
+ "gen_py": attr.bool(),
+ "outs": attr.output_list(),
+ },
+ output_to_genfiles = True,
+ implementation = _proto_gen_impl,
+)
+"""Generates codes from Protocol Buffers definitions.
+
+This rule helps you to implement Skylark macros specific to the target
+language. You should prefer more specific `cc_proto_library `,
+`py_proto_library` and others unless you are adding such wrapper macros.
+
+Args:
+ srcs: Protocol Buffers definition files (.proto) to run the protocol compiler
+ against.
+ deps: a list of dependency labels; must be other proto libraries.
+ includes: a list of include paths to .proto files.
+ protoc: the label of the protocol compiler to generate the sources.
+ plugin: the label of the protocol compiler plugin to be passed to the protocol
+ compiler.
+ plugin_language: the language of the generated sources
+ plugin_options: a list of options to be passed to the plugin
+ gen_cc: generates C++ sources in addition to the ones from the plugin.
+ gen_py: generates Python sources in addition to the ones from the plugin.
+ outs: a list of labels of the expected outputs from the protocol compiler.
+"""
+
+def cc_proto_library(
+ name,
+ srcs = [],
+ deps = [],
+ cc_libs = [],
+ include = None,
+ protoc = "@com_google_protobuf//:protoc",
+ internal_bootstrap_hack = False,
+ use_grpc_plugin = False,
+ default_runtime = "@com_google_protobuf//:protobuf",
+ **kargs):
+ """Bazel rule to create a C++ protobuf library from proto source files
+
+ NOTE: the rule is only an internal workaround to generate protos. The
+ interface may change and the rule may be removed when bazel has introduced
+ the native rule.
+
+ Args:
+ name: the name of the cc_proto_library.
+ srcs: the .proto files of the cc_proto_library.
+ deps: a list of dependency labels; must be cc_proto_library.
+ cc_libs: a list of other cc_library targets depended by the generated
+ cc_library.
+ include: a string indicating the include path of the .proto files.
+ protoc: the label of the protocol compiler to generate the sources.
+ internal_bootstrap_hack: a flag indicate the cc_proto_library is used only
+ for bootstraping. When it is set to True, no files will be generated.
+ The rule will simply be a provider for .proto files, so that other
+ cc_proto_library can depend on it.
+ use_grpc_plugin: a flag to indicate whether to call the grpc C++ plugin
+ when processing the proto files.
+ default_runtime: the implicitly default runtime which will be depended on by
+ the generated cc_library target.
+ **kargs: other keyword arguments that are passed to cc_library.
+
+ """
+
+ includes = []
+ if include != None:
+ includes = [include]
+
+ if internal_bootstrap_hack:
+ # For pre-checked-in generated files, we add the internal_bootstrap_hack
+ # which will skip the codegen action.
+ proto_gen(
+ name = name + "_genproto",
+ srcs = srcs,
+ deps = [s + "_genproto" for s in deps],
+ includes = includes,
+ protoc = protoc,
+ visibility = ["//visibility:public"],
+ )
+
+ # An empty cc_library to make rule dependency consistent.
+ native.cc_library(
+ name = name,
+ **kargs
+ )
+ return
+
+ grpc_cpp_plugin = None
+ if use_grpc_plugin:
+ grpc_cpp_plugin = "//external:grpc_cpp_plugin"
+
+ gen_srcs = _CcSrcs(srcs, use_grpc_plugin)
+ gen_hdrs = _CcHdrs(srcs, use_grpc_plugin)
+ outs = gen_srcs + gen_hdrs
+
+ proto_gen(
+ name = name + "_genproto",
+ srcs = srcs,
+ deps = [s + "_genproto" for s in deps],
+ includes = includes,
+ protoc = protoc,
+ plugin = grpc_cpp_plugin,
+ plugin_language = "grpc",
+ gen_cc = 1,
+ outs = outs,
+ visibility = ["//visibility:public"],
+ )
+
+ if default_runtime and not default_runtime in cc_libs:
+ cc_libs = cc_libs + [default_runtime]
+ if use_grpc_plugin:
+ cc_libs = cc_libs + ["//external:grpc_lib"]
+
+ native.cc_library(
+ name = name,
+ srcs = gen_srcs,
+ hdrs = gen_hdrs,
+ deps = cc_libs + deps,
+ includes = includes,
+ **kargs
+ )
+
+def internal_gen_well_known_protos_java(srcs):
+ """Bazel rule to generate the gen_well_known_protos_java genrule
+
+ Args:
+ srcs: the well known protos
+ """
+ root = Label("%s//protobuf_java" % (REPOSITORY_NAME)).workspace_root
+ pkg = PACKAGE_NAME + "/" if PACKAGE_NAME else ""
+ if root == "":
+ include = " -I%ssrc " % pkg
+ else:
+ include = " -I%s/%ssrc " % (root, pkg)
+ native.genrule(
+ name = "gen_well_known_protos_java",
+ srcs = srcs,
+ outs = [
+ "wellknown.srcjar",
+ ],
+ cmd = "$(location :protoc) --java_out=$(@D)/wellknown.jar" +
+ " %s $(SRCS) " % include +
+ " && mv $(@D)/wellknown.jar $(@D)/wellknown.srcjar",
+ tools = [":protoc"],
+ )
+
+def internal_copied_filegroup(name, srcs, strip_prefix, dest, **kwargs):
+ """Macro to copy files to a different directory and then create a filegroup.
+
+ This is used by the //:protobuf_python py_proto_library target to work around
+ an issue caused by Python source files that are part of the same Python
+ package being in separate directories.
+
+ Args:
+ srcs: The source files to copy and add to the filegroup.
+ strip_prefix: Path to the root of the files to copy.
+ dest: The directory to copy the source files into.
+ **kwargs: extra arguments that will be passesd to the filegroup.
+ """
+ outs = [_RelativeOutputPath(s, strip_prefix, dest) for s in srcs]
+
+ native.genrule(
+ name = name + "_genrule",
+ srcs = srcs,
+ outs = outs,
+ cmd = " && ".join(
+ ["cp $(location %s) $(location %s)" %
+ (s, _RelativeOutputPath(s, strip_prefix, dest)) for s in srcs],
+ ),
+ )
+
+ native.filegroup(
+ name = name,
+ srcs = outs,
+ **kwargs
+ )
+
+def py_proto_library(
+ name,
+ srcs = [],
+ deps = [],
+ py_libs = [],
+ py_extra_srcs = [],
+ include = None,
+ default_runtime = "@com_google_protobuf//:protobuf_python",
+ protoc = "@com_google_protobuf//:protoc",
+ use_grpc_plugin = False,
+ **kargs):
+ """Bazel rule to create a Python protobuf library from proto source files
+
+ NOTE: the rule is only an internal workaround to generate protos. The
+ interface may change and the rule may be removed when bazel has introduced
+ the native rule.
+
+ Args:
+ name: the name of the py_proto_library.
+ srcs: the .proto files of the py_proto_library.
+ deps: a list of dependency labels; must be py_proto_library.
+ py_libs: a list of other py_library targets depended by the generated
+ py_library.
+ py_extra_srcs: extra source files that will be added to the output
+ py_library. This attribute is used for internal bootstrapping.
+ include: a string indicating the include path of the .proto files.
+ default_runtime: the implicitly default runtime which will be depended on by
+ the generated py_library target.
+ protoc: the label of the protocol compiler to generate the sources.
+ use_grpc_plugin: a flag to indicate whether to call the Python C++ plugin
+ when processing the proto files.
+ **kargs: other keyword arguments that are passed to cc_library.
+
+ """
+ outs = _PyOuts(srcs, use_grpc_plugin)
+
+ includes = []
+ if include != None:
+ includes = [include]
+
+ grpc_python_plugin = None
+ if use_grpc_plugin:
+ grpc_python_plugin = "//external:grpc_python_plugin"
+ # Note: Generated grpc code depends on Python grpc module. This dependency
+ # is not explicitly listed in py_libs. Instead, host system is assumed to
+ # have grpc installed.
+
+ proto_gen(
+ name = name + "_genproto",
+ srcs = srcs,
+ deps = [s + "_genproto" for s in deps],
+ includes = includes,
+ protoc = protoc,
+ gen_py = 1,
+ outs = outs,
+ visibility = ["//visibility:public"],
+ plugin = grpc_python_plugin,
+ plugin_language = "grpc",
+ )
+
+ if default_runtime and not default_runtime in py_libs + deps:
+ py_libs = py_libs + [default_runtime]
+
+ native.py_library(
+ name = name,
+ srcs = outs + py_extra_srcs,
+ deps = py_libs + deps,
+ imports = includes,
+ **kargs
+ )
+
+def internal_protobuf_py_tests(
+ name,
+ modules = [],
+ **kargs):
+ """Bazel rules to create batch tests for protobuf internal.
+
+ Args:
+ name: the name of the rule.
+ modules: a list of modules for tests. The macro will create a py_test for
+ each of the parameter with the source "google/protobuf/%s.py"
+ kargs: extra parameters that will be passed into the py_test.
+
+ """
+ for m in modules:
+ s = "python/google/protobuf/internal/%s.py" % m
+ native.py_test(
+ name = "py_%s" % m,
+ srcs = [s],
+ main = s,
+ **kargs
+ )
+
+def check_protobuf_required_bazel_version():
+ """For WORKSPACE files, to check the installed version of bazel.
+
+ This ensures bazel supports our approach to proto_library() depending on a
+ copied filegroup. (Fixed in bazel 0.5.4)
+ """
+ expected = apple_common.dotted_version("0.5.4")
+ current = apple_common.dotted_version(native.bazel_version)
+ if current.compare_to(expected) < 0:
+ fail("Bazel must be newer than 0.5.4")
diff --git a/third_party/systemlibs/syslibs_configure.bzl b/third_party/systemlibs/syslibs_configure.bzl
index dbf4fd6..1b971ec 100644
--- a/third_party/systemlibs/syslibs_configure.bzl
+++ b/third_party/systemlibs/syslibs_configure.bzl
@@ -15,6 +15,8 @@
"boringssl",
"com_github_googleapis_googleapis",
"com_github_googlecloudplatform_google_cloud_cpp",
+ "com_google_protobuf",
+ "com_google_protobuf_cc",
"com_googlesource_code_re2",
"curl",
"cython",
@@ -26,12 +28,14 @@
"icu",
"jpeg",
"jsoncpp_git",
+ "keras_applications_archive",
"lmdb",
"nasm",
"nsync",
"org_sqlite",
"pcre",
"png_archive",
+ "protobuf_archive",
"six_archive",
"snappy",
"swig",
diff --git a/third_party/tflite_mobilenet_float.BUILD b/third_party/tflite_mobilenet_float.BUILD
new file mode 100644
index 0000000..de47ed6
--- /dev/null
+++ b/third_party/tflite_mobilenet_float.BUILD
@@ -0,0 +1,12 @@
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(
+ glob(
+ ["**/*"],
+ exclude = [
+ "BUILD",
+ ],
+ ),
+)
diff --git a/third_party/tflite_mobilenet_quant.BUILD b/third_party/tflite_mobilenet_quant.BUILD
new file mode 100644
index 0000000..de47ed6
--- /dev/null
+++ b/third_party/tflite_mobilenet_quant.BUILD
@@ -0,0 +1,12 @@
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"]) # Apache 2.0
+
+exports_files(
+ glob(
+ ["**/*"],
+ exclude = [
+ "BUILD",
+ ],
+ ),
+)
diff --git a/third_party/toolchains/preconfig/generate/archives.bzl b/third_party/toolchains/preconfig/generate/archives.bzl
new file mode 100644
index 0000000..0850893
--- /dev/null
+++ b/third_party/toolchains/preconfig/generate/archives.bzl
@@ -0,0 +1,27 @@
+load("//tensorflow:version_check.bzl", "parse_bazel_version")
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+def bazel_toolchains_archive():
+ # Not all bazel versions have set native.bazel_version - if it is not set,
+ # fall back to the more compatible version of the toolchains archive.
+ if native.bazel_version and parse_bazel_version(native.bazel_version) >= parse_bazel_version("0.19"):
+ # This version of the toolchains repo is incompatible with older bazel
+ # versions - we can remove this once TensorFlow drops support for bazel
+ # before 0.19.
+ http_archive(
+ name = "bazel_toolchains",
+ sha256 = "41c48a189be489e2d15dec40e0057ea15b95ee5b39cc2a7e6cf663e31432c75e",
+ strip_prefix = "bazel-toolchains-3f8c58fe530fedc446de04673bc1e32985887dea",
+ urls = [
+ "https://github.com/nlopezgi/bazel-toolchains/archive/3f8c58fe530fedc446de04673bc1e32985887dea.tar.gz",
+ ],
+ )
+ else:
+ http_archive(
+ name = "bazel_toolchains",
+ sha256 = "15b5858b1b5541ec44df31b94c3b8672815b31d71215a98398761ea9f4c4eedb",
+ strip_prefix = "bazel-toolchains-6200b238c9c2d137c0d9a7262c80cc71d98e692b",
+ urls = [
+ "https://github.com/bazelbuild/bazel-toolchains/archive/6200b238c9c2d137c0d9a7262c80cc71d98e692b.tar.gz",
+ ],
+ )
diff --git a/third_party/toolchains/preconfig/generate/generate.bzl b/third_party/toolchains/preconfig/generate/generate.bzl
index 2fb3a94..fb2af02 100644
--- a/third_party/toolchains/preconfig/generate/generate.bzl
+++ b/third_party/toolchains/preconfig/generate/generate.bzl
@@ -36,9 +36,7 @@
"TF_NCCL_VERSION": "2",
"CUDNN_INSTALL_PATH": "/usr/lib/x86_64-linux-gnu",
},
- # TODO(klimek): We should use the sources that we currently work on, not
- # just the latest snapshot of tensorflow that is checked in.
- git_repo = "https://github.com/tensorflow/tensorflow",
+ mount_project = "$(mount_project)",
tags = ["manual"],
incompatible_changes_off = True,
)
diff --git a/third_party/toolchains/preconfig/generate/generate.sh b/third_party/toolchains/preconfig/generate/generate.sh
index 37c5211..1f39fcd 100755
--- a/third_party/toolchains/preconfig/generate/generate.sh
+++ b/third_party/toolchains/preconfig/generate/generate.sh
@@ -46,7 +46,7 @@
echo "CUDNN: ${CUDNN_VERSION}"
echo "NCCL: ${NCCL_VERSION}"
-bazel build "${PKG}/generate:${TARGET}"
+bazel build --define=mount_project="${PWD}" "${PKG}/generate:${TARGET}"
cd "${TEMPDIR}"
tar xvf "${ROOT}/bazel-bin/${PKG}/generate/${TARGET}_outputs.tar"
diff --git a/third_party/toolchains/preconfig/win_1803/bazel_018/BUILD b/third_party/toolchains/preconfig/win_1803/bazel_018/BUILD
index c00f005..edd9583 100644
--- a/third_party/toolchains/preconfig/win_1803/bazel_018/BUILD
+++ b/third_party/toolchains/preconfig/win_1803/bazel_018/BUILD
@@ -39,6 +39,9 @@
"x64_windows|msvc-cl": ":cc-compiler-x64_windows",
"x64_windows|msys-gcc": ":cc-compiler-x64_windows_msys",
"x64_windows|mingw-gcc": ":cc-compiler-x64_windows_mingw",
+ "x64_windows_msys": ":cc-compiler-x64_windows_msys",
+ "x64_windows": ":cc-compiler-x64_windows",
+ "armeabi-v7a": ":cc-compiler-armeabi-v7a",
},
)
@@ -54,6 +57,7 @@
static_runtime_libs = [":empty"],
strip_files = ":empty",
supports_param_files = 1,
+ toolchain_identifier = "msys_x64",
)
toolchain(
@@ -83,6 +87,7 @@
static_runtime_libs = [":empty"],
strip_files = ":empty",
supports_param_files = 0,
+ toolchain_identifier = "msys_x64_mingw",
)
toolchain(
@@ -112,6 +117,7 @@
static_runtime_libs = [":empty"],
strip_files = ":empty",
supports_param_files = 1,
+ toolchain_identifier = "msvc_x64",
)
toolchain(
@@ -140,6 +146,7 @@
static_runtime_libs = [":empty"],
strip_files = ":empty",
supports_param_files = 1,
+ toolchain_identifier = "stub_armeabi-v7a",
)
toolchain(
diff --git a/third_party/toolchains/preconfig/win_1803/bazel_018/CROSSTOOL b/third_party/toolchains/preconfig/win_1803/bazel_018/CROSSTOOL
index 04c8bca..38a80c2 100644
--- a/third_party/toolchains/preconfig/win_1803/bazel_018/CROSSTOOL
+++ b/third_party/toolchains/preconfig/win_1803/bazel_018/CROSSTOOL
@@ -14,42 +14,6 @@
major_version: "local"
minor_version: ""
-default_target_cpu: "same_as_host"
-
-default_toolchain {
- cpu: "x64_windows"
- toolchain_identifier: "msvc_x64"
-}
-
-default_toolchain {
- cpu: "local"
- toolchain_identifier: "stub_armeabi-v7a"
-}
-
-default_toolchain {
- cpu: "armeabi-v7a"
- toolchain_identifier: "stub_armeabi-v7a"
-}
-
-default_toolchain {
- cpu: "x64_windows"
- toolchain_identifier: "msvc_x64"
-}
-
-default_toolchain {
- cpu: "x64_windows_msvc"
- toolchain_identifier: "msvc_x64"
-}
-
-default_toolchain {
- cpu: "x64_windows_msys"
- toolchain_identifier: "msys_x64"
-}
-
-default_toolchain {
- cpu: "s390x"
- toolchain_identifier: "msys_x64"
-}
# Android tooling requires a default toolchain for the armeabi-v7a cpu.
toolchain {